pf1e-simulator/tests/test_phase2_candidates.py

"""Phase-2 candidate actions (docs/phase-2-llm-architecture.md P0.3).

``candidates.candidates`` derives a deterministic menu of plausible
``Action`` proposals from the ``Observation`` snapshot's pairs. This module
pins the contract: a normal turn yields a non-empty, capped menu whose kinds
stay within the declared set; the menu is deterministic; no candidate ever
targets an ally (the pairs cover every other combatant, so the side filter is
what keeps the menu hostile-only); a weaponless combatant still gets
movement/defensive options; an enemy in weapon reach gets attack/full_attack;
and harmless spells (cure, buffs) target the caster rather than an enemy.
"""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING

from pf1e_simulator.candidates import _CANDIDATE_KINDS, MAX_CANDIDATES, candidates
from pf1e_simulator.combat import (
    Action,
    CombatantState,
    CombatEngine,
    Observation,
    default_policy,
)
from pf1e_simulator.grid import Grid
from pf1e_simulator.loaders import load_combatant
from pf1e_simulator.map import load_map
from pf1e_simulator.rng import SeededRng
from pf1e_simulator.runner import EncounterSpec, Side, build_states
from pf1e_simulator.spells import SPELLS_DIR, SpellSpec, load_spell_registry

if TYPE_CHECKING:
    from pf1e_simulator.models import Combatant

ROOT = Path(__file__).resolve().parents[1]
MAP = ROOT / "data" / "maps" / "sample_arena.yaml"


class _CaptureAll:
    """Policy wrapper that snapshots every observation, then delegates."""

    def __init__(self) -> None:
        self.observations: dict[str, list[Observation]] = {}

    def policy(self, engine: CombatEngine, state: CombatantState) -> tuple[Action, ...]:
        obs = engine._build_observation(state)
        self.observations.setdefault(state.combatant.id, []).append(obs)
        return default_policy(engine, state)


def _c(path: str) -> Combatant:
    return load_combatant(ROOT / "data" / path)


def _representative_spec() -> EncounterSpec:
    """Caster-heavy 3v4 on the sample arena (walls, rubble, water, pillars)."""
    return EncounterSpec(
        map=load_map(MAP),
        round_cap=20,
        sides=(
            Side(
                name="players",
                combatants=(_c("pcs/oni.json"), _c("pcs/misty.json"), _c("pcs/nairda.json")),
                zone="A",
            ),
            Side(
                name="monsters",
                combatants=(
                    _c("monsters/lich.json"),
                    _c("monsters/thrune_agent.json"),
                    _c("monsters/orc.json"),
                    _c("monsters/goblin.json"),
                ),
                zone="B",
            ),
        ),
    )


def _weaponless_spec() -> EncounterSpec:
    """1v1: a goblin with no usable weapons against an orc."""
    goblin = _c("monsters/goblin.json").model_copy(update={"attacks": ()})
    return EncounterSpec(
        map=load_map(MAP),
        round_cap=20,
        sides=(
            Side(name="players", combatants=(goblin,), zone="A"),
            Side(name="monsters", combatants=(_c("monsters/orc.json"),), zone="B"),
        ),
    )


def _run_observations(
    spec: EncounterSpec, seed: int, registry: dict[str, SpellSpec]
) -> dict[str, list[Observation]]:
    grid = Grid.from_spec(spec.map)
    states = build_states(spec)
    cap = _CaptureAll()
    CombatEngine(
        SeededRng(seed),
        grid,
        states,
        round_cap=spec.round_cap,
        policy=cap.policy,
        spell_registry=registry,
    ).run()
    return cap.observations


def _pair_target_side(pair: dict[str, object], obs: Observation) -> object:
    """Side of the pair's target combatant, or None when out of range."""
    index = pair.get("target_index")
    if not isinstance(index, int) or not (0 <= index < len(obs.combatants)):
        return None
    return obs.combatants[index].get("side")


def test_candidates_normal_turn_non_empty_and_supported_kinds() -> None:
    registry = load_spell_registry(SPELLS_DIR)
    observations = _run_observations(_representative_spec(), 1, registry)
    obs = next(obs_list[0] for obs_list in observations.values())
    menu = candidates(obs)
    assert menu
    assert len(menu) <= MAX_CANDIDATES
    kinds = {a.kind for a in menu}
    assert kinds <= _CANDIDATE_KINDS
    assert {"move", "total_defense", "delay"} <= kinds


def test_candidates_deterministic() -> None:
    registry = load_spell_registry(SPELLS_DIR)
    observations = _run_observations(_representative_spec(), 1, registry)
    obs = next(obs_list[0] for obs_list in observations.values())
    assert candidates(obs) == candidates(obs)


def test_candidates_no_allies_targeted() -> None:
    registry = load_spell_registry(SPELLS_DIR)
    observations = _run_observations(_representative_spec(), 1, registry)
    for obs_list in observations.values():
        for obs in obs_list:
            acting = obs.combatants[obs.acting_combatant]
            acting_side = acting["side"]
            acting_id = acting["id"]
            for action in candidates(obs):
                if action.target_id is None or action.target_id == acting_id:
                    continue
                target = next(
                    (c for c in obs.combatants if c["id"] == action.target_id), None
                )
                assert target is not None
                assert target["side"] != acting_side


def test_candidates_weaponless_gets_movement_and_defense() -> None:
    registry = load_spell_registry(SPELLS_DIR)
    observations = _run_observations(_weaponless_spec(), 1, registry)
    obs = observations["goblin"][0]
    kinds = {a.kind for a in candidates(obs)}
    assert {"move", "total_defense", "ready", "delay"} <= kinds
    assert not ({"attack", "full_attack", "charge"} & kinds)


def test_candidates_in_range_gets_attack() -> None:
    registry = load_spell_registry(SPELLS_DIR)
    observations = _run_observations(_representative_spec(), 1, registry)
    for obs_list in observations.values():
        for obs in obs_list:
            acting = obs.combatants[obs.acting_combatant]
            acting_side = acting["side"]
            enemy_pairs = [
                p for p in obs.pairs if _pair_target_side(p, obs) != acting_side
            ]
            if not any(p.get("weapon_for") is not None for p in enemy_pairs):
                continue
            kinds = {a.kind for a in candidates(obs)}
            assert "attack" in kinds
            assert "full_attack" in kinds
            return
    msg = "no observation with an enemy in weapon reach found"
    raise AssertionError(msg)


def test_candidates_harmless_spells_target_self() -> None:
    registry = load_spell_registry(SPELLS_DIR)
    observations = _run_observations(_representative_spec(), 1, registry)
    obs = observations["oni"][0]
    acting = obs.combatants[obs.acting_combatant]
    acting_id = acting["id"]
    harmless_casts = [
        a
        for a in candidates(obs)
        if a.kind == "cast_spell"
        and a.spell_name in {"Guidance", "Shield of Faith", "Stabilize"}
    ]
    assert harmless_casts
    assert all(a.target_id == acting_id for a in harmless_casts)