"""Package U8 — final validation: an observation-only LLM policy drives the
Phase-2 capabilities (U1-U7b) end-to-end through the real engine.

Each test wires a ``llm_policy`` (``LLMPolicy(observation)``) that
reads ONLY the frozen Observation snapshot — never ``engine.*`` — and emits
one newly-unblocked tactical decision per capability: ``ready`` (U4),
``delay`` (U3), ``total_defense`` (U5a), ``move`` with an explicit
``destination`` (U5c), ``attack`` with an explicit ``weapon_id`` (U5d), and
the U1/U2 feedback loop (``last_turn_results`` + ``initiative_order``).
The last test chains ready -> total_defense -> weapon-selected attack over
three rounds. Every roll (initiative included) is scripted with
``ScriptedRng``; a too-short or too-long queue fails the test, pinning the
exact roll count of each scenario.
"""
from __future__ import annotations

from collections.abc import Callable

from pf1e_simulator.combat import (
    Action,
    ActionResult,
    CombatantState,
    CombatEngine,
    LLMPolicy,
    Observation,
)
from pf1e_simulator.dice import parse_dice
from pf1e_simulator.models import (
    AbilityScores,
    ACProfile,
    AttackSpec,
    Combatant,
    DamageComponent,
    Saves,
)
from pf1e_simulator.rng import ScriptedRng
from tests.test_observation_enriched import make_grid, make_state

_WAIT: tuple[Action, ...] = (Action(kind="wait"),)

Decide = Callable[[str, int, Observation], tuple[Action, ...]]


def _sword(cid: str) -> AttackSpec:
    return AttackSpec(
        id=f"{cid}-sword",
        name="short sword",
        kind="melee",
        attack_bonus=5,
        damage=[DamageComponent(formula=parse_dice("1d6"), types=["slashing"])],
    )


def _bow(cid: str) -> AttackSpec:
    return AttackSpec(
        id=f"{cid}-bow",
        name="short bow",
        kind="ranged",
        attack_bonus=4,
        damage=[DamageComponent(formula=parse_dice("1d6"), types=["piercing"])],
        range_increment_ft=60,
    )


def _combatant(
    id_: str,
    attacks: list[AttackSpec],
    *,
    initiative_mod: int = 0,
) -> Combatant:
    return Combatant(
        id=id_,
        name=id_,
        level=1,
        size="Medium",
        abilities=AbilityScores(
            str_score=10, dex_score=10, con_score=10,
            int_score=10, wis_score=10, cha_score=10,
        ),
        hp_max=100,
        ac=ACProfile(total=10, touch=10, flat_footed=10),
        bab=1,
        initiative_mod=initiative_mod,
        speed_land_ft=30,
        saves=Saves(fort=0, ref=0, will=0),
        attacks=attacks,
    )


def _engine(
    rng: list[int],
    states: list[CombatantState],
    policy: LLMPolicy,
    round_cap: int,
) -> CombatEngine:
    return CombatEngine(
        ScriptedRng(rng),
        make_grid(),
        states,
        round_cap=round_cap,
        llm_policy=policy,
    )


def _me(obs: Observation) -> dict[str, object]:
    return obs.combatants[obs.acting_combatant]


def _enemy_id(obs: Observation) -> str:
    """First active enemy id, read purely from the observation pairs."""
    me = _me(obs)
    for pair in obs.pairs:
        target = obs.combatants[pair["target_index"]]
        if target["side"] != me["side"] and target["active"]:
            return str(target["id"])
    msg = "no enemy pair in the observation"
    raise AssertionError(msg)


def _counting(decide: Decide) -> LLMPolicy:
    """LLMPolicy calling ``decide(id, call_no, obs)`` per turn (obs only)."""
    calls: dict[str, int] = {}

    def policy(obs: Observation) -> tuple[Action, ...]:
        cid = str(_me(obs)["id"])
        n = calls.get(cid, 0)
        calls[cid] = n + 1
        return decide(cid, n, obs)

    return policy


def _order(transcript: tuple[str, ...], *needles: str) -> tuple[int, ...]:
    """Indices of the first transcript line containing each needle."""
    return tuple(
        next(i for i, line in enumerate(transcript) if needle in line)
        for needle in needles
    )


# ── ready (U4): a readied attack fires off-turn when the foe moves ────────────


def test_llm_policy_ready_fires_when_enemy_moves() -> None:
    """Given: an llm_policy that readies an attack on the enemy when it
    enters reach, and an enemy policy that walks in
    When: the battle runs (initiative 10/5, readied swing d20 10, dmg 3)
    Then: the readied attack fires off-turn during the enemy's move, the
    readied action is consumed, and the enemy took the damage."""

    def decide(cid: str, _n: int, obs: Observation) -> tuple[Action, ...]:
        if cid == "hero":
            return (
                Action(
                    kind="ready",
                    ready_trigger="enemy_enters_reach",
                    target_id=_enemy_id(obs),
                ),
            )
        return (Action(kind="move", target_id=_enemy_id(obs)),)

    hero = make_state(_combatant("hero", [_sword("hero")], initiative_mod=2), "players", (0, 0))
    foe = make_state(_combatant("foe", [_sword("foe")]), "monsters", (0, 2))
    result = _engine([10, 5, 10, 3], [hero, foe], _counting(decide), 1).run()
    t = result.transcript

    assert "round 1 hero: ready (enemy_enters_reach)" in t
    assert "round 1 hero: ready fires (enemy_enters_reach)" in t
    assert (
        "round 1 hero: short sword vs foe d20=10+5=15 AC 10 -> HIT 3 damage (100->97)"
        in t
    )
    idx = _order(
        t,
        "round 1 hero: ready fires (enemy_enters_reach)",
        "round 1 foe: move (0,2)->(0,1)",
    )
    assert idx == tuple(sorted(idx))
    assert hero.readied is None
    assert hero.readied_trigger is None
    assert foe.hp == 97


# ── delay (U3): the reordered initiative shows up in later observations ───────


def test_llm_policy_delay_reorders_initiative_and_reports_feedback() -> None:
    """Given: an llm_policy that delays on hero's first turn, then waits
    When: the battle runs (initiative hero 7 / foe 5, no other rolls)
    Then: hero acts AFTER the foe in the same round, the foe's and hero's
    re-queued observations show initiative_order ("foe", "hero"), and the
    re-queued turn's last_turn_results carry the ok delay feedback."""
    orders: list[tuple[str, ...]] = []
    feedback: list[object] = []

    def decide(cid: str, n: int, obs: Observation) -> tuple[Action, ...]:
        orders.append(obs.initiative_order)
        feedback.append(_me(obs)["last_turn_results"])
        if cid == "hero" and n == 0:
            return (Action(kind="delay"),)
        return _WAIT

    hero = make_state(_combatant("hero", [_sword("hero")], initiative_mod=2), "players", (0, 0))
    foe = make_state(_combatant("foe", [_sword("foe")], initiative_mod=1), "monsters", (0, 3))
    result = _engine([5, 4], [hero, foe], _counting(decide), 1).run()
    t = result.transcript

    assert "round 1 hero: delay" in t
    idx = _order(
        t,
        "round 1 hero: delay",
        "round 1 foe: wait",
        "round 1 hero: wait",
    )
    assert idx == tuple(sorted(idx))
    # hero turn 1, foe turn, hero re-queued turn — exactly three policy calls.
    assert orders == [("hero", "foe"), ("foe", "hero"), ("foe", "hero")]
    assert feedback[2] == (
        {"kind": "delay", "ok": True, "reason": "delayed"},
    )


# ── total defense (U5a): the AC bonus is readable and applied ─────────────────


def test_llm_policy_total_defense_raises_ac_seen_by_enemy_policy() -> None:
    """Given: an llm_policy where hero takes total defense and the foe
    attacks after reading hero's snapshot
    When: the battle runs (initiative 12/5, foe swing d20 8 -> 13)
    Then: the foe's observation exposes defensive_ac_bonus 4 on hero and the
    same roll that would hit AC 10 misses AC 14."""
    seen_bonus: list[object] = []

    def decide(cid: str, _n: int, obs: Observation) -> tuple[Action, ...]:
        if cid == "hero":
            return (Action(kind="total_defense"),)
        hero_c = next(c for c in obs.combatants if c["id"] == "hero")
        seen_bonus.append(hero_c["defensive_ac_bonus"])
        return (Action(kind="attack", target_id=_enemy_id(obs)),)

    hero = make_state(_combatant("hero", [_sword("hero")], initiative_mod=2), "players", (0, 0))
    foe = make_state(_combatant("foe", [_sword("foe")]), "monsters", (0, 1))
    result = _engine([10, 5, 8], [hero, foe], _counting(decide), 1).run()
    t = result.transcript

    assert seen_bonus == [4]
    assert "round 1 hero: total defense (+4 AC)" in t
    assert "round 1 foe: short sword vs hero d20=8+5=13 AC 14 -> MISS" in t
    assert hero.hp == 100


# ── movement destination (U5c): destination picked from the observation ───────


def test_llm_policy_moves_to_first_free_observed_cell() -> None:
    """Given: an llm_policy that scans obs.terrain and the combatants' pos
    to pick the first free floor cell, then moves there
    When: the battle runs (initiative 12/5, no other rolls)
    Then: hero's pos is exactly that cell and the move reports ok."""
    chosen: list[object] = []

    def decide(cid: str, _n: int, obs: Observation) -> tuple[Action, ...]:
        if cid != "hero":
            return _WAIT
        me = _me(obs)
        occupied = {c["pos"] for c in obs.combatants if c["active"]}
        for r, row in enumerate(obs.terrain):
            for col, kind in enumerate(row):
                cell = (r, col)
                if kind == "floor" and cell != me["pos"] and cell not in occupied:
                    chosen.append(cell)
                    return (Action(kind="move", destination=cell),)
        return _WAIT

    hero = make_state(_combatant("hero", [_sword("hero")], initiative_mod=2), "players", (0, 0))
    foe = make_state(_combatant("foe", [_sword("foe")]), "monsters", (0, 4))
    engine = _engine([10, 5], [hero, foe], _counting(decide), 1)
    result = engine.run()

    assert chosen == [(0, 1)]
    assert hero.pos == (0, 1)
    assert "round 1 hero: move -> (0,1)" in result.transcript
    assert hero.last_action_results == [
        ActionResult(kind="move", ok=True, reason="ok")
    ]


# ── weapon selection (U5d): the policy names the bow, the bow swings ─────────


def test_llm_policy_attacks_with_the_chosen_bow() -> None:
    """Given: a hero with sword + bow adjacent to a non-threatening foe, and
    an llm_policy that picks the ranged weapon id from the observation
    When: the battle runs (initiative 12/5, shot d20 10 -> 14, dmg 3)
    Then: the transcript shows the short bow swinging (not the sword) and
    the auto-selection (sword, first usable) was overridden."""

    def decide(cid: str, _n: int, obs: Observation) -> tuple[Action, ...]:
        if cid != "hero":
            return _WAIT
        bow_id = next(str(a["id"]) for a in _me(obs)["attacks"] if a["kind"] == "ranged")
        return (Action(kind="attack", target_id=_enemy_id(obs), weapon_id=bow_id),)

    hero = make_state(
        _combatant("hero", [_sword("hero"), _bow("hero")], initiative_mod=2),
        "players",
        (0, 0),
    )
    foe = make_state(_combatant("foe", [_bow("foe")]), "monsters", (0, 1))
    result = _engine([10, 5, 10, 3], [hero, foe], _counting(decide), 1).run()
    t = result.transcript

    assert "round 1 hero: short bow vs foe d20=10+4=14 AC 10 -> HIT 3 damage (100->97)" in t
    assert not any("short sword vs foe" in line for line in t)
    assert foe.hp == 97
    assert hero.last_action_results == [
        ActionResult(kind="attack", ok=True, reason="ok")
    ]


# ── feedback loop (U1+U2): correct a rejected action from last_turn_results ──


def test_llm_policy_fixes_rejected_move_from_last_turn_results() -> None:
    """Given: an llm_policy that moves with no target on round 1 (rejected
    missing_target) and re-reads its own last_turn_results each round
    When: the battle runs two rounds (initiative 12/5, no other rolls)
    Then: the round-2 observation carries the round-1 rejection, the policy
    switches to total_defense accordingly, and initiative_order is present
    and stable on both of hero's turns."""
    seen_orders: list[tuple[str, ...]] = []
    seen_results: list[object] = []

    def decide(cid: str, _n: int, obs: Observation) -> tuple[Action, ...]:
        if cid != "hero":
            return _WAIT
        seen_orders.append(obs.initiative_order)
        results = _me(obs)["last_turn_results"]
        seen_results.append(results)
        if results and any(not bool(item["ok"]) for item in results):
            return (Action(kind="total_defense"),)
        return (Action(kind="move"),)

    hero = make_state(_combatant("hero", [_sword("hero")], initiative_mod=2), "players", (0, 0))
    foe = make_state(_combatant("foe", [_sword("foe")]), "monsters", (0, 4))
    result = _engine([10, 5], [hero, foe], _counting(decide), 2).run()
    t = result.transcript

    assert seen_orders == [("hero", "foe"), ("hero", "foe")]
    assert seen_results[0] == ()
    assert seen_results[1] == (
        {"kind": "move", "ok": False, "reason": "missing_target"},
    )
    assert "round 1 hero: move rejected (missing_target)" in t
    assert "round 2 hero: total defense (+4 AC)" in t


# ── capstone: ready -> total_defense -> weapon-selected attack over 3 rounds ─


def test_llm_policy_multi_round_tactical_sequence() -> None:
    """Given: hero (sword + bow, init 12) vs foe (sword, init 5) and an
    llm_policy playing round 1 ready, round 2 total_defense, round 3 a
    bow attack chosen by weapon_id from the observation; the foe moves in
    round 1, attacks in round 2 and waits in round 3
    When: the battle runs (8 scripted rolls: 2 initiative, readied swing +
    dmg, foe swing, foe AoO, bow swing + dmg)
    Then: every capability lands in the transcript in order, the readied
    action is consumed, and hp ends hero 100 / foe 94."""

    def decide(cid: str, n: int, obs: Observation) -> tuple[Action, ...]:
        enemy = _enemy_id(obs)
        if cid == "hero":
            if n == 0:
                return (
                    Action(
                        kind="ready",
                        ready_trigger="enemy_enters_reach",
                        target_id=enemy,
                    ),
                )
            if n == 1:
                return (Action(kind="total_defense"),)
            bow_id = next(
                str(a["id"]) for a in _me(obs)["attacks"] if a["kind"] == "ranged"
            )
            return (Action(kind="attack", target_id=enemy, weapon_id=bow_id),)
        if n == 0:
            return (Action(kind="move", target_id=enemy),)
        if n == 1:
            return (Action(kind="attack", target_id=enemy),)
        return _WAIT

    hero = make_state(
        _combatant("hero", [_sword("hero"), _bow("hero")], initiative_mod=2),
        "players",
        (0, 0),
    )
    foe = make_state(_combatant("foe", [_sword("foe")]), "monsters", (0, 2))
    engine = _engine(
        # init hero 10+2, foe 5; readied swing 10+5 HIT dmg 3; foe swing
        # 8+5=13 vs AC 14 MISS; foe AoO 1+5=6 vs AC 10 MISS; bow 10+4 HIT dmg 3.
        [10, 5, 10, 3, 8, 1, 10, 3],
        [hero, foe],
        _counting(decide),
        3,
    )
    result = engine.run()
    t = result.transcript

    assert "round 1 hero: ready fires (enemy_enters_reach)" in t
    assert (
        "round 1 hero: short sword vs foe d20=10+5=15 AC 10 -> HIT 3 damage (100->97)"
        in t
    )
    assert "round 2 hero: total defense (+4 AC)" in t
    assert "round 2 foe: short sword vs hero d20=8+5=13 AC 14 -> MISS" in t
    assert "round 3 foe: AoO vs hero d20=1+5=6 AC 10 -> MISS" in t
    assert (
        "round 3 hero: short bow vs foe d20=10+4=14 AC 10 -> HIT 3 damage (97->94)"
        in t
    )
    idx = _order(
        t,
        "round 1 hero: ready (enemy_enters_reach)",
        "round 1 hero: ready fires (enemy_enters_reach)",
        "round 2 hero: total defense",
        "round 3 hero: short bow vs foe",
    )
    assert idx == tuple(sorted(idx))
    assert hero.readied is None
    assert hero.hp == 100
    assert foe.hp == 94