pf1e-simulator/tests/test_retrieve_weapon.py

"""Tests for the ``retrieve_weapon`` action (recovering a stolen weapon).

L7 of the ``llm-policy-contract`` plan: ``Action`` gains a
``retrieve_weapon`` kind and a ``weapon_id`` field so a policy can restore a
weapon rendered unusable in ``state.weapon_stolen``. The weapon is removed
from ``state.weapon_stolen`` and becomes usable again via ``usable_attacks()``.
(The steal maneuver no longer fills this set: APG Steal forbids taking wielded
weapons.)

No default policy emits ``retrieve_weapon``; it is only accepted by
``_execute`` (usage via custom policy / LLM).
"""

from __future__ import annotations

from pf1e_simulator.combat import Action, CombatantState, CombatEngine
from pf1e_simulator.grid import Grid
from pf1e_simulator.map import MapSpec, TerrainType
from pf1e_simulator.models import (
    AbilityScores,
    ACProfile,
    AttackSpec,
    Combatant,
    DamageComponent,
    Saves,
)
from pf1e_simulator.rng import ScriptedRng, SeededRng


def make_combatant(cid: str, *, attack_bonus: int = 2) -> Combatant:
    attack = AttackSpec(
        id=f"{cid}-w",
        name="sword",
        kind="melee",
        attack_bonus=attack_bonus,
        damage=[DamageComponent(formula="1d4", types=["slashing"])],
    )
    return Combatant(
        id=cid,
        name=cid,
        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=6,
        ac=ACProfile(total=13, touch=13, flat_footed=13),
        bab=attack_bonus,
        initiative_mod=0,
        speed_land_ft=30,
        saves=Saves(fort=0, ref=0, will=0),
        attacks=[attack],
    )


def make_state(combatant: Combatant, side: str, pos: tuple[int, int]) -> CombatantState:
    return CombatantState(combatant=combatant, side=side, pos=pos, hp=combatant.hp_max)


def make_grid() -> Grid:
    legend = {".": TerrainType(type="floor", move_cost=1)}
    spec = MapSpec(name="test", terrain=tuple(["." * 8] * 8), legend=legend)
    return Grid.from_spec(spec)


def make_engine(queue: list[int], states: list[CombatantState]) -> CombatEngine:
    return CombatEngine(ScriptedRng(queue), make_grid(), states)


def _weapon_ids(state: CombatantState) -> list[str]:
    return [w.id for w in state.usable_attacks()]


# --- Test 1: retrieve restores a stolen weapon -------------------------------

def test_retrieve_restores_stolen_weapon() -> None:
    a_spec = make_combatant("a", attack_bonus=2)
    b_spec = make_combatant("b")
    a = make_state(a_spec, "players", (1, 1))
    b = make_state(b_spec, "monsters", (1, 2))
    # The steal maneuver no longer takes wielded weapons (APG Steal), so the
    # stolen state is set directly — retrieve_weapon is the mechanism under
    # test here.
    b.weapon_stolen.add("b-w")
    engine = make_engine([], [a, b])
    assert "b-w" not in _weapon_ids(b)
    engine._execute(b, Action(kind="retrieve_weapon", weapon_id="b-w"))
    assert not b.weapon_stolen
    assert "b-w" in _weapon_ids(b)


# --- Test 2: weapon_id=None -> nothing to retrieve, state unchanged ----------

def test_retrieve_none_logs_nothing_to_retrieve() -> None:
    a_spec = make_combatant("a", attack_bonus=2)
    b_spec = make_combatant("b")
    a = make_state(a_spec, "players", (1, 1))
    b = make_state(b_spec, "monsters", (1, 2))
    b.weapon_stolen.add("b-w")
    engine = make_engine([], [a, b])
    engine._execute(b, Action(kind="retrieve_weapon"))
    assert "nothing to retrieve" in engine._transcript[-1]
    assert b.weapon_stolen == {"b-w"}


# --- Test 3: unknown id (never stolen) -> log, state unchanged ---------------

def test_retrieve_unknown_id_logs_nothing_to_retrieve() -> None:
    a_spec = make_combatant("a", attack_bonus=2)
    b_spec = make_combatant("b")
    a = make_state(a_spec, "players", (1, 1))
    b = make_state(b_spec, "monsters", (1, 2))
    b.weapon_stolen.add("b-w")
    engine = make_engine([], [a, b])
    engine._execute(b, Action(kind="retrieve_weapon", weapon_id="ghost"))
    assert "nothing to retrieve" in engine._transcript[-1]
    assert b.weapon_stolen == {"b-w"}


# --- Test 4: steal leaves wielded weapons; retrieve restores a stolen id -----

def test_steal_then_retrieve_restores_weapon() -> None:
    # APG Steal: a wielded weapon cannot be taken — the maneuver provokes,
    # succeeds, and steals nothing; the owner keeps swinging. A weapon that IS
    # stolen (model state set directly) is restored by retrieve_weapon.
    a_spec = make_combatant("a", attack_bonus=2)
    b_spec = make_combatant("b")
    a = make_state(a_spec, "players", (1, 1))
    b = make_state(b_spec, "monsters", (1, 2))
    engine = make_engine([1, 12], [a, b])
    engine._execute(a, Action(kind="maneuver", target_id="b", maneuver="steal"))
    assert not b.weapon_stolen
    assert engine.weapon_for(b, a) is not None
    b.weapon_stolen.add("b-w")
    engine._execute(b, Action(kind="retrieve_weapon", weapon_id="b-w"))
    assert not b.weapon_stolen
    assert engine.weapon_for(b, a) is not None
    assert "retrieve b-w (restored)" in engine._transcript[-1]


# --- Test 5: determinism -----------------------------------------------------

def _run_with_retrieve() -> tuple[str, ...]:
    a_spec = make_combatant("a", attack_bonus=5)
    b_spec = make_combatant("b")
    a = make_state(a_spec, "players", (1, 1))
    b = make_state(b_spec, "monsters", (1, 2))
    b.weapon_stolen.add("b-w")

    def policy(engine: CombatEngine, state: CombatantState) -> tuple[Action, ...]:
        if state.combatant.id == "b":
            return (Action(kind="retrieve_weapon", weapon_id="b-w"),)
        return (Action(kind="wait"),)

    engine = CombatEngine(SeededRng(42), make_grid(), [a, b], policy=policy)
    return engine.run().transcript


def test_determinism_same_seed_same_transcript() -> None:
    assert _run_with_retrieve() == _run_with_retrieve()


# --- Test 6: free draw combined with a move at BAB +1 ------------------------


def _move_retrieve_attack_run(bab: int) -> CombatEngine:
    a_spec = make_combatant("a", attack_bonus=bab)
    b_spec = make_combatant("b")
    a = make_state(a_spec, "players", (1, 1))
    b = make_state(b_spec, "monsters", (1, 4))
    a.weapon_stolen.add("a-w")

    def policy(engine: CombatEngine, state: CombatantState) -> tuple[Action, ...]:
        if state.combatant.id == "a":
            return (
                Action(kind="move", target_id="b"),
                Action(kind="retrieve_weapon", weapon_id="a-w"),
                Action(kind="attack", target_id="b"),
            )
        return (Action(kind="wait"),)

    engine = CombatEngine(
        ScriptedRng([10, 5, 10, 3]), make_grid(), [a, b], policy=policy, round_cap=1
    )
    engine.run()
    return engine


def test_bab_plus1_draws_for_free_and_still_attacks() -> None:
    # SRD "Move Actions": "If you have a base attack bonus of +1 or higher,
    # you may draw a weapon as a free action combined with a regular move."
    # Given: a BAB +2 combatant whose turn sequence is (move, retrieve, attack)
    # When: the turn runs
    # Then: all three actions execute — the draw rides the move for free and
    # the standard attack still happens.
    engine = _move_retrieve_attack_run(bab=2)
    a = engine._states[0]
    assert not a.weapon_stolen
    kinds = [(r.kind, r.ok) for r in a.last_action_results]
    assert ("move", True) in kinds
    assert ("retrieve_weapon", True) in kinds
    assert ("attack", True) in kinds
    assert any(" vs b " in line for line in engine._transcript)


def test_bab_zero_draw_still_costs_the_standard_slot() -> None:
    # Given: the same sequence but BAB 0 — the draw is a move action, which
    # converts the standard slot, leaving nothing for the attack.
    # Then: the attack is downgraded out.
    engine = _move_retrieve_attack_run(bab=0)
    a = engine._states[0]
    assert not a.weapon_stolen
    assert any(
        r.kind == "attack" and r.reason == "downgraded:action_economy"
        for r in a.last_action_results
    )