"""Tests for atm.state_machine."""

from __future__ import annotations

import pytest

from atm.state_machine import DotColor, State, StateMachine, Transition

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


def _drive(sm: StateMachine, events: list[tuple[DotColor, float]]) -> list[Transition]:
    return [sm.feed(color, ts) for color, ts in events]


def _buy_sequence_to_primed(sm: StateMachine, start_ts: float = 1.0) -> None:
    """Drive sm from IDLE → PRIMED_BUY (does not fire)."""
    sm.feed("turquoise", start_ts)       # → ARMED_BUY
    sm.feed("dark_green", start_ts + 1)  # → PRIMED_BUY


def _sell_sequence_to_primed(sm: StateMachine, start_ts: float = 1.0) -> None:
    """Drive sm from IDLE → PRIMED_SELL (does not fire)."""
    sm.feed("yellow", start_ts)          # → ARMED_SELL
    sm.feed("dark_red", start_ts + 1)    # → PRIMED_SELL


# ---------------------------------------------------------------------------
# 1. clean_buy
# ---------------------------------------------------------------------------


def test_clean_buy() -> None:
    sm = StateMachine()

    t1 = sm.feed("turquoise", 1.0)
    assert t1.prev == State.IDLE
    assert t1.next == State.ARMED_BUY
    assert t1.reason == "arm"
    assert sm.state == State.ARMED_BUY

    t2 = sm.feed("gray", 2.0)
    assert t2.prev == State.ARMED_BUY
    assert t2.next == State.ARMED_BUY
    assert t2.reason == "persist"
    assert sm.state == State.ARMED_BUY

    t3 = sm.feed("dark_green", 3.0)
    assert t3.prev == State.ARMED_BUY
    assert t3.next == State.PRIMED_BUY
    assert t3.reason == "prime"
    assert sm.state == State.PRIMED_BUY

    t4 = sm.feed("light_green", 4.0)
    assert t4.prev == State.PRIMED_BUY
    assert t4.next == State.IDLE
    assert t4.reason == "fire"
    assert t4.trigger == "BUY"
    assert t4.locked is False
    assert sm.state == State.IDLE


# ---------------------------------------------------------------------------
# 2. clean_sell (mirror of clean_buy)
# ---------------------------------------------------------------------------


def test_clean_sell() -> None:
    sm = StateMachine()

    t1 = sm.feed("yellow", 1.0)
    assert t1.prev == State.IDLE
    assert t1.next == State.ARMED_SELL
    assert t1.reason == "arm"

    t2 = sm.feed("gray", 2.0)
    assert t2.prev == State.ARMED_SELL
    assert t2.next == State.ARMED_SELL
    assert t2.reason == "persist"

    t3 = sm.feed("dark_red", 3.0)
    assert t3.prev == State.ARMED_SELL
    assert t3.next == State.PRIMED_SELL
    assert t3.reason == "prime"

    t4 = sm.feed("light_red", 4.0)
    assert t4.prev == State.PRIMED_SELL
    assert t4.next == State.IDLE
    assert t4.reason == "fire"
    assert t4.trigger == "SELL"
    assert t4.locked is False


# ---------------------------------------------------------------------------
# 3. cooled
# ---------------------------------------------------------------------------


def test_cooled() -> None:
    sm = StateMachine()
    _buy_sequence_to_primed(sm, start_ts=1.0)
    assert sm.state == State.PRIMED_BUY

    t = sm.feed("gray", 10.0)
    assert t.prev == State.PRIMED_BUY
    assert t.next == State.IDLE
    assert t.reason == "cooled"
    assert t.trigger is None


# ---------------------------------------------------------------------------
# 4. opposite_rearm from ARMED_BUY
# ---------------------------------------------------------------------------


def test_opposite_rearm_from_armed_buy() -> None:
    sm = StateMachine()
    sm.feed("turquoise", 1.0)  # → ARMED_BUY
    assert sm.state == State.ARMED_BUY

    t = sm.feed("yellow", 2.0)
    assert t.prev == State.ARMED_BUY
    assert t.next == State.ARMED_SELL
    assert t.reason == "opposite_rearm"
    assert sm.state == State.ARMED_SELL


# ---------------------------------------------------------------------------
# 5. opposite_rearm from PRIMED_BUY
# ---------------------------------------------------------------------------


def test_opposite_rearm_from_primed_buy() -> None:
    sm = StateMachine()
    _buy_sequence_to_primed(sm, start_ts=1.0)
    assert sm.state == State.PRIMED_BUY

    t = sm.feed("yellow", 5.0)
    assert t.prev == State.PRIMED_BUY
    assert t.next == State.ARMED_SELL
    assert t.reason == "opposite_rearm"
    assert sm.state == State.ARMED_SELL


# ---------------------------------------------------------------------------
# 6. lockout_same_direction
# ---------------------------------------------------------------------------


def test_lockout_same_direction() -> None:
    sm = StateMachine(lockout_s=240)

    # First fire at t=100
    _buy_sequence_to_primed(sm, start_ts=90.0)
    t_fire1 = sm.feed("light_green", 100.0)
    assert t_fire1.trigger == "BUY"
    assert t_fire1.locked is False

    # Re-prime
    _buy_sequence_to_primed(sm, start_ts=110.0)

    # Second fire at t=200 — inside lockout window (200 - 100 = 100 < 240)
    t_fire2 = sm.feed("light_green", 200.0)
    assert t_fire2.trigger == "BUY"
    assert t_fire2.locked is True
    assert t_fire2.next == State.IDLE

    # Re-prime again
    _buy_sequence_to_primed(sm, start_ts=300.0)

    # Third fire at t=341 — outside lockout (341 - 200 = 141 < 240, still locked)
    # Actually 341 - 100 would be outside but last fire is at t=200:
    # 341 - 200 = 141 < 240 → still locked
    # We need t > 200 + 240 = 440
    _buy_sequence_to_primed(sm, start_ts=430.0)
    t_fire3 = sm.feed("light_green", 441.0)
    assert t_fire3.trigger == "BUY"
    assert t_fire3.locked is False


def test_lockout_same_direction_boundary() -> None:
    """Spec requirement: fire BUY @ t=100; fire again @ t=341 → locked=False (241 >= 240)."""
    sm = StateMachine(lockout_s=240)

    _buy_sequence_to_primed(sm, start_ts=90.0)
    t_fire1 = sm.feed("light_green", 100.0)
    assert t_fire1.locked is False

    # Re-prime and fire just inside window: 339-100=239 < 240 → locked
    _buy_sequence_to_primed(sm, start_ts=110.0)
    t_locked = sm.feed("light_green", 339.0)
    assert t_locked.locked is True

    # last_fire is now 339. Re-prime and fire just outside: 580-339=241 >= 240 → unlocked
    _buy_sequence_to_primed(sm, start_ts=340.0)
    t_free = sm.feed("light_green", 580.0)
    assert t_free.locked is False


# ---------------------------------------------------------------------------
# 7. lockout_does_not_block_opposite
# ---------------------------------------------------------------------------


def test_lockout_does_not_block_opposite() -> None:
    sm = StateMachine(lockout_s=240)

    # Fire BUY at t=100
    _buy_sequence_to_primed(sm, start_ts=90.0)
    sm.feed("light_green", 100.0)

    # Drive SELL sequence — opposite direction must not be locked
    _sell_sequence_to_primed(sm, start_ts=110.0)
    t_sell = sm.feed("light_red", 200.0)
    assert t_sell.trigger == "SELL"
    assert t_sell.locked is False


# ---------------------------------------------------------------------------
# 8. phase_skip from ARMED_BUY (light_green without priming)
# ---------------------------------------------------------------------------


def test_phase_skip_armed_buy() -> None:
    sm = StateMachine()
    sm.feed("turquoise", 1.0)  # → ARMED_BUY
    assert sm.state == State.ARMED_BUY

    t = sm.feed("light_green", 2.0)
    assert t.prev == State.ARMED_BUY
    assert t.next == State.IDLE
    assert t.reason == "phase_skip"
    assert t.trigger is None


# ---------------------------------------------------------------------------
# 9. noise_from_idle
# ---------------------------------------------------------------------------


def test_noise_from_idle() -> None:
    sm = StateMachine()
    t = sm.feed("dark_green", 1.0)
    assert t.prev == State.IDLE
    assert t.next == State.IDLE
    assert t.reason == "noise"
    assert t.trigger is None
    assert sm.state == State.IDLE


# ---------------------------------------------------------------------------
# 10. refresh_arm_ts
# ---------------------------------------------------------------------------


def test_refresh_arm_ts() -> None:
    sm = StateMachine()
    sm.feed("turquoise", 1.0)  # arm at t=1

    t1 = sm.feed("turquoise", 5.0)  # refresh at t=5
    assert t1.prev == State.ARMED_BUY
    assert t1.next == State.ARMED_BUY
    assert t1.reason == "refresh"
    assert t1.arm_ts == 5.0

    t2 = sm.feed("turquoise", 9.0)  # refresh again at t=9
    assert t2.arm_ts == 9.0


# ---------------------------------------------------------------------------
# 11. fired_in_session — public API for catchup suppression after fire
# ---------------------------------------------------------------------------


def test_fired_in_session_fresh_fsm() -> None:
    """Fresh FSM — neither direction has fired."""
    sm = StateMachine()
    assert sm.fired_in_session("BUY") is False
    assert sm.fired_in_session("SELL") is False


def test_fired_in_session_after_buy_fire() -> None:
    """After a BUY fire: BUY=True, SELL=False (direction-scoped)."""
    sm = StateMachine(lockout_s=240)
    sm.feed("turquoise", 1.0)
    sm.feed("dark_green", 2.0)
    t = sm.feed("light_green", 3.0)
    assert t.reason == "fire"

    assert sm.fired_in_session("BUY") is True
    assert sm.fired_in_session("SELL") is False


def test_fired_in_session_after_sell_fire() -> None:
    """Mirror — after SELL fire: SELL=True, BUY=False."""
    sm = StateMachine(lockout_s=240)
    sm.feed("yellow", 1.0)
    sm.feed("dark_red", 2.0)
    t = sm.feed("light_red", 3.0)
    assert t.reason == "fire"

    assert sm.fired_in_session("SELL") is True
    assert sm.fired_in_session("BUY") is False


def test_fired_in_session_both_directions() -> None:
    """Fire both directions — both True."""
    sm = StateMachine(lockout_s=240)
    # BUY cycle
    sm.feed("turquoise", 1.0)
    sm.feed("dark_green", 2.0)
    sm.feed("light_green", 3.0)
    # SELL cycle
    sm.feed("yellow", 100.0)
    sm.feed("dark_red", 101.0)
    sm.feed("light_red", 102.0)

    assert sm.fired_in_session("BUY") is True
    assert sm.fired_in_session("SELL") is True


# ---------------------------------------------------------------------------
# 11. exhaustive — parameterize over every (state, color) pair
# ---------------------------------------------------------------------------

ALL_STATES = list(State)
ALL_COLORS: list[DotColor] = [
    "turquoise", "yellow", "dark_green", "dark_red",
    "light_green", "light_red", "gray",
]

FIRE_DIRECTIONS: dict[str, str] = {
    State.PRIMED_BUY.value: "BUY",
    State.PRIMED_SELL.value: "SELL",
}

VALID_STATES = set(State)


def _sm_in_state(target: State) -> StateMachine:
    """Return a fresh StateMachine already in the given state."""
    sm = StateMachine()
    match target:
        case State.IDLE:
            pass
        case State.ARMED_BUY:
            sm.feed("turquoise", 1.0)
        case State.ARMED_SELL:
            sm.feed("yellow", 1.0)
        case State.PRIMED_BUY:
            sm.feed("turquoise", 1.0)
            sm.feed("dark_green", 2.0)
        case State.PRIMED_SELL:
            sm.feed("yellow", 1.0)
            sm.feed("dark_red", 2.0)
    assert sm.state == target, f"Setup failed: wanted {target}, got {sm.state}"
    return sm


@pytest.mark.parametrize("state", ALL_STATES)
@pytest.mark.parametrize("color", ALL_COLORS)
def test_exhaustive(state: State, color: DotColor) -> None:
    sm = _sm_in_state(state)
    t = sm.feed(color, 10.0)

    # (a) resulting state is valid
    assert t.next in VALID_STATES, f"Invalid next state: {t.next}"

    # (b) reason is non-empty
    assert t.reason, f"Empty reason for ({state}, {color})"

    # (c) if fire, trigger matches direction
    if t.reason == "fire":
        expected_dir = FIRE_DIRECTIONS.get(state.value)
        assert t.trigger == expected_dir, (
            f"Wrong trigger for fire from {state}: got {t.trigger}, expected {expected_dir}"
        )
