"""ATM trading state machine — pure stdlib, no cv2/numpy."""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Literal


class State(Enum):
    IDLE = "IDLE"
    ARMED_BUY = "ARMED_BUY"
    PRIMED_BUY = "PRIMED_BUY"
    ARMED_SELL = "ARMED_SELL"
    PRIMED_SELL = "PRIMED_SELL"


DotColor = Literal[
    "turquoise", "yellow", "dark_green", "dark_red", "light_green", "light_red", "gray"
]


@dataclass(frozen=True)
class Transition:
    prev: State
    next: State
    event: str
    reason: str
    trigger: Literal["BUY", "SELL"] | None = None
    locked: bool = False
    arm_ts: float | None = None
    prime_ts: float | None = None
    fire_ts: float | None = None


class StateMachine:
    """Finite state machine for ATM dot-color trading signals.

    BUY side and SELL side are mirrors of each other:
      turquoise ↔ yellow
      dark_green ↔ dark_red
      light_green ↔ light_red

    Lockout: after a fire event, same-direction fires within lockout_s
    seconds produce locked=True transitions (caller suppresses notify).
    Opposite direction is unaffected.
    """

    def __init__(self, lockout_s: int = 240) -> None:
        self._state = State.IDLE
        self._lockout_s = lockout_s
        # arm_ts / prime_ts are reset per-arm / per-prime cycle
        self._arm_ts: float | None = None
        self._prime_ts: float | None = None
        # last fire timestamps per direction
        self._last_fire: dict[str, float] = {}

    @property
    def state(self) -> State:
        return self._state

    def fired_in_session(self, direction: str) -> bool:
        """True if this FSM instance has recorded any FIRE in the given direction.

        Used by catchup logic to suppress synth-arm after a cycle already fired.
        Residual dark_green / dark_red dots for the remainder of the 15m window
        look identical to startup-catchup from (color, state) alone, but must
        be classified as noise, not a new prime signal.
        """
        return direction in self._last_fire

    # ------------------------------------------------------------------
    # Core feed
    # ------------------------------------------------------------------

    def feed(self, color: DotColor, ts: float) -> Transition:  # noqa: C901
        prev = self._state

        match self._state:
            case State.IDLE:
                t = self._from_idle(color, ts)
            case State.ARMED_BUY:
                t = self._from_armed(color, ts, direction="BUY")
            case State.ARMED_SELL:
                t = self._from_armed(color, ts, direction="SELL")
            case State.PRIMED_BUY:
                t = self._from_primed(color, ts, direction="BUY")
            case State.PRIMED_SELL:
                t = self._from_primed(color, ts, direction="SELL")
            case _:  # pragma: no cover
                t = self._noise(color, ts)

        assert t.prev == prev
        self._state = t.next
        return t

    # ------------------------------------------------------------------
    # State handlers
    # ------------------------------------------------------------------

    def _from_idle(self, color: DotColor, ts: float) -> Transition:
        prev = State.IDLE
        match color:
            case "turquoise":
                self._arm_ts = ts
                self._prime_ts = None
                return Transition(prev, State.ARMED_BUY, color, "arm", arm_ts=ts)
            case "yellow":
                self._arm_ts = ts
                self._prime_ts = None
                return Transition(prev, State.ARMED_SELL, color, "arm", arm_ts=ts)
            case "light_green" | "light_red":
                return Transition(prev, State.IDLE, color, "phase_skip")
            case _:
                # dark_green, dark_red, gray → noise
                return Transition(prev, State.IDLE, color, "noise")

    def _from_armed(
        self, color: DotColor, ts: float, direction: Literal["BUY", "SELL"]
    ) -> Transition:
        """Handle events from ARMED_BUY or ARMED_SELL."""
        prev = State.ARMED_BUY if direction == "BUY" else State.ARMED_SELL
        # Map color names to logical roles for the given direction
        arm_color = "turquoise" if direction == "BUY" else "yellow"
        opposite_arm = "yellow" if direction == "BUY" else "turquoise"
        prime_color = "dark_green" if direction == "BUY" else "dark_red"
        skip_correct = "light_green" if direction == "BUY" else "light_red"
        skip_wrong = "light_red" if direction == "BUY" else "light_green"
        primed_state = State.PRIMED_BUY if direction == "BUY" else State.PRIMED_SELL
        opposite_state = State.ARMED_SELL if direction == "BUY" else State.ARMED_BUY

        match color:
            case c if c == "gray":
                return Transition(prev, prev, color, "persist", arm_ts=self._arm_ts)
            case c if c == arm_color:
                # refresh: update arm_ts
                self._arm_ts = ts
                return Transition(prev, prev, color, "refresh", arm_ts=ts)
            case c if c == prime_color:
                self._prime_ts = ts
                return Transition(
                    prev, primed_state, color, "prime",
                    arm_ts=self._arm_ts, prime_ts=ts,
                )
            case c if c == opposite_arm:
                self._arm_ts = ts
                self._prime_ts = None
                return Transition(
                    prev, opposite_state, color, "opposite_rearm", arm_ts=ts
                )
            case c if c == skip_correct:
                # phase_skip — no fire
                self._arm_ts = None
                self._prime_ts = None
                return Transition(prev, State.IDLE, color, "phase_skip")
            case c if c == skip_wrong:
                self._arm_ts = None
                self._prime_ts = None
                return Transition(prev, State.IDLE, color, "phase_skip")
            case _:
                # dark_red (for BUY) or dark_green (for SELL) → ignore
                return Transition(prev, prev, color, "ignore", arm_ts=self._arm_ts)

    def _from_primed(
        self, color: DotColor, ts: float, direction: Literal["BUY", "SELL"]
    ) -> Transition:
        """Handle events from PRIMED_BUY or PRIMED_SELL."""
        prev = State.PRIMED_BUY if direction == "BUY" else State.PRIMED_SELL
        accumulate_color = "dark_green" if direction == "BUY" else "dark_red"
        ignore_color = "dark_red" if direction == "BUY" else "dark_green"
        fire_color = "light_green" if direction == "BUY" else "light_red"
        skip_color = "light_red" if direction == "BUY" else "light_green"
        arm_color = "turquoise" if direction == "BUY" else "yellow"
        opposite_arm = "yellow" if direction == "BUY" else "turquoise"
        armed_self = State.ARMED_BUY if direction == "BUY" else State.ARMED_SELL
        armed_opp = State.ARMED_SELL if direction == "BUY" else State.ARMED_BUY

        match color:
            case c if c == accumulate_color:
                return Transition(
                    prev, prev, color, "accumulate",
                    arm_ts=self._arm_ts, prime_ts=self._prime_ts,
                )
            case c if c == ignore_color:
                return Transition(
                    prev, prev, color, "ignore",
                    arm_ts=self._arm_ts, prime_ts=self._prime_ts,
                )
            case c if c == fire_color:
                locked = self._is_locked(direction, ts)
                fire_ts = ts
                self._last_fire[direction] = ts
                self._arm_ts = None
                self._prime_ts = None
                return Transition(
                    prev, State.IDLE, color, "fire",
                    trigger=direction,
                    locked=locked,
                    fire_ts=fire_ts,
                )
            case c if c == skip_color:
                self._arm_ts = None
                self._prime_ts = None
                return Transition(prev, State.IDLE, color, "phase_skip")
            case c if c == "gray":
                self._arm_ts = None
                self._prime_ts = None
                return Transition(prev, State.IDLE, color, "cooled")
            case c if c == arm_color:
                self._arm_ts = ts
                self._prime_ts = None
                return Transition(prev, armed_self, color, "rearm", arm_ts=ts)
            case c if c == opposite_arm:
                self._arm_ts = ts
                self._prime_ts = None
                return Transition(prev, armed_opp, color, "opposite_rearm", arm_ts=ts)
            case _:
                return Transition(
                    prev, prev, color, "noise",
                    arm_ts=self._arm_ts, prime_ts=self._prime_ts,
                )

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _noise(self, color: DotColor, ts: float) -> Transition:  # pragma: no cover
        return Transition(self._state, self._state, color, "noise")

    def _is_locked(self, direction: str, ts: float) -> bool:
        last = self._last_fire.get(direction)
        if last is None:
            return False
        return (ts - last) < self._lockout_s

    # ------------------------------------------------------------------
    # Public lockout API — used by fire_on_phase_skip handler outside the
    # FSM. Mirrors _is_locked / _last_fire without leaking private attrs.
    # ------------------------------------------------------------------

    def is_locked(self, direction: str, ts: float) -> bool:
        """True if a FIRE in `direction` at ts would be within the lockout window."""
        return self._is_locked(direction, ts)

    def record_fire(self, direction: str, ts: float) -> None:
        """Mark a FIRE for `direction` at ts, starting the lockout timer.

        Used by backstop handlers (e.g. fire_on_phase_skip) that emit a
        fire-equivalent alert without going through the natural FSM path.
        """
        self._last_fire[direction] = ts
