'''
Deterministic offline market-data replay backend.

'''
from __future__ import annotations

from contextlib import (
    asynccontextmanager as acm,
)
from dataclasses import dataclass
from datetime import datetime
import os
from pathlib import Path
from typing import (
    Any,
    Literal,
)

import msgspec
import numpy as np
import pendulum
import tractor
import trio
from trio_typing import TaskStatus

from piker.accounting import MktPair
from piker.brokers import (
    DataUnavailable,
    SymbolNotFound,
)
from piker.data._source import def_iohlcv_fields
from piker.data.validate import FeedInit
from piker.types import Struct


name: str = 'replay'
_scenario_env: str = 'PIKER_REPLAY_SCENARIO'


class ReplayBar(Struct, frozen=True):
    '''
    One normalized OHLCV fixture row.

    '''
    index: int
    time: int
    open: float
    high: float
    low: float
    close: float
    volume: float


class ReplayTick(Struct, frozen=True):
    '''
    One normalized tick in a replay quote.

    '''
    type: str
    price: float
    size: float


class ReplayQuote(Struct, frozen=True):
    '''
    One sequenced provider event from a replay scenario.

    '''
    sequence: int
    fqme: str
    broker_ts: int
    last: float
    ticks: list[ReplayTick]


class ReplayFailure(Struct, frozen=True):
    '''
    A failure which can be armed at one event sequence.

    '''
    sequence: int
    code: str
    message: str


class ReplayScenario(Struct, frozen=True):
    '''
    Versioned input for one deterministic replay generation.

    '''
    version: int
    scenario_id: str
    markets: list[MktPair]
    history_1s: list[ReplayBar]
    history_1m: list[ReplayBar]
    quotes: list[ReplayQuote]
    failures: list[ReplayFailure] = []


ReplayState = Literal[
    'ready',
    'paused',
    'failed',
    'exhausted',
]


class Advance(Struct, frozen=True, tag=True):
    '''
    Publish one exact replay event.

    '''
    command_id: int
    event_sequence: int


class Pause(Struct, frozen=True, tag=True):
    '''
    Close the replay producer gate.

    '''
    command_id: int


class Resume(Struct, frozen=True, tag=True):
    '''
    Open the replay producer gate.

    '''
    command_id: int


class FailAt(Struct, frozen=True, tag=True):
    '''
    Arm one scenario-declared failure injection.

    '''
    command_id: int
    event_sequence: int


class AwaitState(Struct, frozen=True, tag=True):
    '''
    Wait for one provider state without scheduler polling.

    '''
    command_id: int
    state: ReplayState


class Snapshot(Struct, frozen=True, tag=True):
    '''
    Request the current deterministic replay state.

    '''
    command_id: int


ReplayCommand = (
    Advance
    | Pause
    | Resume
    | FailAt
    | AwaitState
    | Snapshot
)


class ReplayRecord(Struct, frozen=True):
    '''
    One deterministic control transcript entry.

    '''
    command_id: int
    command: str
    state: ReplayState
    event_sequence: int
    ok: bool
    error: str = ''


class ReplaySnapshot(Struct, frozen=True, tag=True):
    '''
    Provider state returned through the control protocol.

    '''
    scenario_id: str
    state: ReplayState
    command_id: int
    event_sequence: int
    subscriber_active: bool
    armed_failure: int | None
    transcript: list[ReplayRecord]


class ReplayAck(Struct, frozen=True, tag=True):
    '''
    Correlated result for one replay control command.

    '''
    command_id: int
    ok: bool
    snapshot: ReplaySnapshot
    error: str = ''


ReplayPayload = (
    ReplayCommand
    | ReplayAck
    | ReplaySnapshot
)


# hmm maybe put this in a new piker._testing.brokers ?
# or is there some reason to put it alongside the production
# backends?
def load_scenario(
    path: Path | str,
) -> ReplayScenario:
    '''
    Decode and validate one versioned replay scenario.

    '''
    scenario_path: Path = Path(path)
    scenario: ReplayScenario = msgspec.json.decode(
        scenario_path.read_bytes(),
        type=ReplayScenario,
    )
    if scenario.version != 1:
        raise ValueError(
            f'Unsupported replay scenario version: '
            f'{scenario.version}'
        )
    if not scenario.markets:
        raise ValueError('Replay scenario has no markets')
    if not scenario.quotes:
        raise ValueError('Replay scenario has no quotes')

    quote_sequences: list[int] = [
        quote.sequence
        for quote in scenario.quotes
    ]
    expected: list[int] = list(range(
        quote_sequences[0],
        quote_sequences[-1] + 1,
    ))
    if quote_sequences != expected:
        raise ValueError(
            'Replay quote sequences must be contiguous'
        )

    market_fqmes: set[str] = {
        mkt.fqme
        for mkt in scenario.markets
    }
    unknown_fqmes: set[str] = {
        quote.fqme
        for quote in scenario.quotes
    } - market_fqmes
    if unknown_fqmes:
        raise ValueError(
            f'Replay quotes reference unknown markets: '
            f'{sorted(unknown_fqmes)!r}'
        )
    return scenario


def _scenario_path() -> Path:
    path: str | None = os.environ.get(_scenario_env)
    if not path:
        raise RuntimeError(
            f'Set `{_scenario_env}` to a replay scenario path'
        )
    return Path(path)


def _quote_to_msg(
    quote: ReplayQuote,
) -> dict[str, Any]:
    return {
        'symbol': quote.fqme,
        'last': quote.last,
        'broker_ts': quote.broker_ts,
        'brokerd_ts': quote.broker_ts,
        'replay_seq': quote.sequence,
        'ticks': [
            {
                'type': tick.type,
                'price': tick.price,
                'size': tick.size,
            }
            for tick in quote.ticks
        ],
    }


def _bars_to_array(
    bars: list[ReplayBar],
) -> np.ndarray:
    return np.asarray(
        [
            (
                bar.index,
                bar.time,
                bar.open,
                bar.high,
                bar.low,
                bar.close,
                bar.volume,
            )
            for bar in bars
        ],
        dtype=def_iohlcv_fields,
    )


def _find_market(
    scenario: ReplayScenario,
    fqme: str,
) -> MktPair:
    normalized: str = fqme.lower()
    for mkt in scenario.markets:
        aliases: set[str] = {
            mkt.fqme.lower(),
            mkt.bs_fqme.lower(),
            mkt.bs_mktid.lower(),
        }
        if normalized in aliases:
            return mkt
    raise SymbolNotFound(fqme)


async def get_mkt_info(
    fqme: str,
) -> tuple[MktPair, MktPair]:
    '''
    Return normalized and backend market records from the fixture.

    '''
    scenario: ReplayScenario = load_scenario(_scenario_path())
    mkt: MktPair = _find_market(scenario, fqme)
    return mkt, mkt


@acm
async def open_history_client(
    mkt: MktPair,
):
    '''
    Open a bounded in-memory OHLCV history client.

    '''
    scenario: ReplayScenario = load_scenario(_scenario_path())
    _find_market(scenario, mkt.fqme)
    frames: dict[int, np.ndarray] = {
        1: _bars_to_array(scenario.history_1s),
        60: _bars_to_array(scenario.history_1m),
    }

    async def get_ohlc(
        timeframe: float,
        end_dt: datetime | None = None,
        start_dt: datetime | None = None,

    ) -> tuple[np.ndarray, datetime, datetime]:
        period: int = int(timeframe)
        try:
            frame: np.ndarray = frames[period]
        except KeyError:
            raise DataUnavailable(
                f'Unsupported replay timeframe: {timeframe}'
            ) from None

        selected: np.ndarray = frame
        if start_dt is not None:
            start_ts: float = start_dt.timestamp()
            selected = selected[selected['time'] >= start_ts]
        if end_dt is not None:
            end_ts: float = end_dt.timestamp()
            selected = selected[selected['time'] < end_ts]
        if not selected.size:
            raise DataUnavailable(
                f'No replay history before {end_dt!r}'
            )

        result: np.ndarray = selected.copy()
        start: datetime = pendulum.from_timestamp(
            int(result['time'][0]),
        )
        end: datetime = pendulum.from_timestamp(
            int(result['time'][-1]),
        )
        return result, start, end

    yield get_ohlc, {
        'erlangs': 1,
        'rate': 1,
    }


@dataclass
class _RuntimeRequest:
    command: ReplayCommand
    reply: trio.MemorySendChannel[ReplayAck]


class _ReplayRuntime:
    '''
    Actor-local owner of replay progression and transcript state.

    '''
    def __init__(
        self,
        scenario: ReplayScenario,
    ) -> None:
        self.scenario = scenario
        self.command_tx: (
            trio.MemorySendChannel[_RuntimeRequest]
            | None
        ) = None
        self.command_id: int = 0
        self.event_sequence: int = scenario.quotes[0].sequence
        self.producer_paused: bool = False
        self.subscriber_active: bool = True
        self.armed_failure: int | None = None
        self.failed: bool = False
        self.transcript: list[ReplayRecord] = []
        self._state_changed: trio.Event = trio.Event()

    @property
    def state(self) -> ReplayState:
        if self.failed:
            return 'failed'
        if (
            self.event_sequence
            == self.scenario.quotes[-1].sequence
        ):
            return 'exhausted'
        if (
            self.producer_paused
            or not self.subscriber_active
        ):
            return 'paused'
        return 'ready'

    def snapshot(self) -> ReplaySnapshot:
        return ReplaySnapshot(
            scenario_id=self.scenario.scenario_id,
            state=self.state,
            command_id=self.command_id,
            event_sequence=self.event_sequence,
            subscriber_active=self.subscriber_active,
            armed_failure=self.armed_failure,
            transcript=list(self.transcript),
        )

    def signal_state_change(self) -> None:
        changed: trio.Event = self._state_changed
        self._state_changed = trio.Event()
        changed.set()

    async def await_state(
        self,
        state: ReplayState,
    ) -> None:
        while self.state != state:
            changed: trio.Event = self._state_changed
            await changed.wait()

    def record(
        self,
        command: ReplayCommand,
        ok: bool,
        error: str = '',
    ) -> ReplayAck:
        self.transcript.append(ReplayRecord(
            command_id=command.command_id,
            command=type(command).__name__,
            state=self.state,
            event_sequence=self.event_sequence,
            ok=ok,
            error=error,
        ))
        return ReplayAck(
            command_id=command.command_id,
            ok=ok,
            error=error,
            snapshot=self.snapshot(),
        )


_runtime: _ReplayRuntime | None = None
_runtime_path: Path | None = None


def _get_runtime() -> _ReplayRuntime:
    global _runtime, _runtime_path
    path: Path = _scenario_path().resolve()
    if (
        _runtime is None
        or _runtime_path != path
    ):
        _runtime = _ReplayRuntime(load_scenario(path))
        _runtime_path = path
    return _runtime


def on_feed_subscription_change(
    bs_fqme: str,
    active: bool,
) -> None:
    '''
    Synchronize the provider gate with feed-bus subscriptions.

    '''
    runtime: _ReplayRuntime = _get_runtime()
    _find_market(runtime.scenario, bs_fqme)
    if runtime.subscriber_active != active:
        runtime.subscriber_active = active
        runtime.signal_state_change()


async def _execute_command(
    runtime: _ReplayRuntime,
    command: ReplayCommand,
    send_chan: trio.abc.SendChannel,
) -> ReplayAck:
    expected_id: int = runtime.command_id + 1
    if command.command_id != expected_id:
        error: str = (
            f'Expected command_id={expected_id}, got '
            f'{command.command_id}'
        )
        return runtime.record(command, False, error)

    runtime.command_id = command.command_id
    match command:
        case Pause():
            runtime.producer_paused = True
            runtime.signal_state_change()

        case Resume():
            runtime.producer_paused = False
            runtime.signal_state_change()

        case AwaitState(state=state):
            await runtime.await_state(state)

        case FailAt(event_sequence=sequence):
            failures: dict[int, ReplayFailure] = {
                failure.sequence: failure
                for failure in runtime.scenario.failures
            }
            if sequence not in failures:
                error = (
                    f'No scenario failure at event_sequence='
                    f'{sequence}'
                )
                return runtime.record(command, False, error)
            runtime.armed_failure = sequence

        case Advance(event_sequence=sequence):
            if runtime.state != 'ready':
                error = (
                    f'Cannot advance replay while state='
                    f'{runtime.state!r}'
                )
                return runtime.record(command, False, error)

            expected_sequence: int = runtime.event_sequence + 1
            if sequence != expected_sequence:
                error = (
                    f'Expected event_sequence={expected_sequence}, '
                    f'got {sequence}'
                )
                return runtime.record(command, False, error)

            if runtime.armed_failure == sequence:
                failure: ReplayFailure = next(
                    failure
                    for failure in runtime.scenario.failures
                    if failure.sequence == sequence
                )
                runtime.failed = True
                runtime.signal_state_change()
                error = f'{failure.code}: {failure.message}'
                return runtime.record(command, False, error)

            quote: ReplayQuote = runtime.scenario.quotes[
                sequence
                - runtime.scenario.quotes[0].sequence
            ]
            mkt: MktPair = _find_market(
                runtime.scenario,
                quote.fqme,
            )
            await send_chan.send({
                mkt.bs_fqme: _quote_to_msg(quote),
            })
            runtime.event_sequence = sequence
            runtime.signal_state_change()

        case Snapshot():
            pass

    return runtime.record(command, True)


async def stream_quotes(
    send_chan: trio.abc.SendChannel,
    symbols: list[str],
    feed_is_live: trio.Event,
    loglevel: str | None = None,
    task_status: TaskStatus[
        tuple[list[FeedInit], dict[str, Any]]
    ] = trio.TASK_STATUS_IGNORED,
) -> None:
    '''
    Publish fixture quotes only after acknowledged control commands.

    '''
    if len(symbols) != 1:
        raise ValueError(
            'The replay backend currently supports one symbol'
        )

    runtime: _ReplayRuntime = _get_runtime()
    mkt: MktPair = _find_market(runtime.scenario, symbols[0])
    first: ReplayQuote = runtime.scenario.quotes[0]
    if first.fqme != mkt.fqme:
        raise ValueError(
            f'First replay quote does not target {mkt.fqme!r}'
        )

    command_tx: trio.MemorySendChannel[_RuntimeRequest]
    command_rx: trio.MemoryReceiveChannel[_RuntimeRequest]
    command_tx, command_rx = trio.open_memory_channel(0)
    runtime.command_tx = command_tx
    try:
        async with (
            send_chan,
            command_rx,
        ):
            task_status.started((
                [FeedInit(mkt_info=mkt)],
                _quote_to_msg(first),
            ))
            feed_is_live.set()

            async for request in command_rx:
                ack: ReplayAck = await _execute_command(
                    runtime,
                    request.command,
                    send_chan,
                )
                async with request.reply:
                    await request.reply.send(ack)
    finally:
        runtime.command_tx = None


@tractor.context(pld_spec=ReplayPayload)
async def open_replay_control(
    ctx: tractor.Context,
) -> None:
    '''
    Open the typed replay command and acknowledgement stream.

    '''
    runtime: _ReplayRuntime = _get_runtime()
    if runtime.command_tx is None:
        raise RuntimeError('Replay quote stream is not running')

    await ctx.started(runtime.snapshot())
    async with ctx.open_stream() as stream:
        with ctx.pld_rx.limit_plds(spec=ReplayCommand):
            command: ReplayCommand
            async for command in stream:
                reply_tx: trio.MemorySendChannel[ReplayAck]
                reply_rx: trio.MemoryReceiveChannel[ReplayAck]
                reply_tx, reply_rx = trio.open_memory_channel(1)
                request = _RuntimeRequest(
                    command=command,
                    reply=reply_tx,
                )
                await runtime.command_tx.send(request)
                async with reply_rx:
                    ack: ReplayAck = await reply_rx.receive()
                await stream.send(ack)


@tractor.context
async def open_symbol_search(
    ctx: tractor.Context,
) -> None:
    '''
    Search replay fixture markets by normalized FQME substring.

    '''
    scenario: ReplayScenario = load_scenario(_scenario_path())
    await ctx.started()
    async with ctx.open_stream() as stream:
        pattern: str
        async for pattern in stream:
            lowered: str = pattern.lower()
            matches: dict[str, dict[str, Any]] = {
                mkt.fqme: mkt.to_dict()
                for mkt in scenario.markets
                if lowered in mkt.fqme.lower()
            }
            await stream.send(matches)


_datad_mods: list[str] = []
__enable_modules__: list[str] = []
