pf1e-simulator/tests/test_phase2_replay.py

"""Phase-2 replay log (docs/phase-2-llm-architecture.md P0.5).

``replay.py`` records every ``(observation, action)`` pair a policy produced,
keyed by the execution-order ``call_seq``, and serializes the whole battle
(spec + seed + round cap + calls) to JSONL. ``replay()`` turns the log back
into an observation-only policy returning the recorded actions.

These tests pin the contract: JSONL round-trips to canonically-equal records
(``observation_key`` ignores JSON's tuple/list collapse), a replayed battle is
bit-identical to the recorded one (including from a serialized log), a
duplicate ``call_seq`` is rejected, and playback fails loudly on a mismatched
observation and once the log is exhausted.
"""

from __future__ import annotations

from dataclasses import replace
from pathlib import Path
from typing import TYPE_CHECKING

import pytest

from pf1e_simulator.candidates import candidates
from pf1e_simulator.loaders import load_combatant
from pf1e_simulator.map import load_map
from pf1e_simulator.replay import (
    ReplayError,
    ReplayExhaustedError,
    ReplayLog,
    ReplayMismatchError,
    ReplayRecord,
    observation_key,
    replay,
)
from pf1e_simulator.runner import EncounterSpec, Side, run_encounter
from pf1e_simulator.spells import SPELLS_DIR, load_spell_registry

if TYPE_CHECKING:
    from pf1e_simulator.combat import Action, CombatantState, Observation
    from pf1e_simulator.models import Combatant
    from pf1e_simulator.runner import EncounterReport

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


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


def _duel_spec() -> EncounterSpec:
    """A fast 1v1 (goblin archer vs orc) on the sample arena."""
    return EncounterSpec(
        map=load_map(MAP),
        round_cap=20,
        sides=(
            Side(name="players", combatants=(_c("monsters/goblin.json"),), zone="A"),
            Side(name="monsters", combatants=(_c("monsters/orc.json"),), zone="B"),
        ),
    )


class _Recorder:
    """Legacy two-argument policy: log every call, then pick ``candidates[0]``."""

    def __init__(self, spec: EncounterSpec, seed: int, round_cap: int) -> None:
        self.log = ReplayLog(spec=spec, seed=seed, round_cap=round_cap)

    def policy(
        self, observation: Observation, _state: CombatantState
    ) -> tuple[Action, ...]:
        action = candidates(observation)[0]
        self.log.add(
            ReplayRecord(
                call_seq=len(self.log),
                observation=observation,
                action=action,
            )
        )
        return (action,)


def _record(spec: EncounterSpec, *, seed: int = SEED) -> tuple[_Recorder, EncounterReport]:
    """Record one battle and return the recorder plus its captured report."""
    registry = load_spell_registry(SPELLS_DIR)
    recorder = _Recorder(spec, seed=seed, round_cap=spec.round_cap)
    report = run_encounter(
        spec,
        runs=1,
        seed=seed,
        llm_policy=recorder.policy,
        spell_registry=registry,
        capture_first_transcript=True,
    )
    return recorder, report


def test_jsonl_round_trip_preserves_records() -> None:
    """A log serializes to JSONL and loads back to structurally-equal records."""
    spec = _duel_spec()
    recorder, _report = _record(spec)
    log = recorder.log

    assert len(log) > 0
    assert [record.call_seq for record in log.records] == list(range(len(log)))

    text = log.to_jsonl()
    restored = ReplayLog.from_jsonl(text)

    assert len(restored) == len(log)
    assert [record.call_seq for record in restored.records] == [
        record.call_seq for record in log.records
    ]
    assert [
        observation_key(record.observation) for record in restored.records
    ] == [observation_key(record.observation) for record in log.records]
    assert [record.action for record in restored.records] == [
        record.action for record in log.records
    ]
    assert restored.spec == log.spec
    assert restored.seed == log.seed
    assert restored.round_cap == log.round_cap
    assert restored.to_jsonl() == text


def test_replay_reproduces_the_recorded_battle() -> None:
    """Replaying the in-memory log matches the recorded battle exactly."""
    spec = _duel_spec()
    registry = load_spell_registry(SPELLS_DIR)
    recorder, recorded = _record(spec)

    replayed = run_encounter(
        spec,
        runs=1,
        seed=SEED,
        llm_policy=replay(recorder.log),
        spell_registry=registry,
        capture_first_transcript=True,
    )

    assert replayed.wins == recorded.wins
    assert replayed.draws == recorded.draws
    assert replayed.avg_rounds == recorded.avg_rounds
    assert replayed.transcript == recorded.transcript


def test_replay_from_serialized_log_reproduces_the_battle() -> None:
    """A log thawed from JSONL replays the battle bit-for-bit."""
    spec = _duel_spec()
    registry = load_spell_registry(SPELLS_DIR)
    recorder, recorded = _record(spec)

    restored = ReplayLog.from_jsonl(recorder.log.to_jsonl())
    replayed = run_encounter(
        spec,
        runs=1,
        seed=SEED,
        llm_policy=replay(restored),
        spell_registry=registry,
        capture_first_transcript=True,
    )

    assert replayed.wins == recorded.wins
    assert replayed.draws == recorded.draws
    assert replayed.avg_rounds == recorded.avg_rounds
    assert replayed.transcript == recorded.transcript


def test_replay_rejects_a_mismatched_observation() -> None:
    """An observation that differs from the recorded one raises a mismatch error."""
    spec = _duel_spec()
    recorder, _report = _record(spec)
    log = recorder.log
    policy = replay(log)
    first = log.records[0].observation
    tampered = replace(first, round=first.round + 1)

    with pytest.raises(ReplayMismatchError):
        policy(tampered)


def test_replay_returns_each_recorded_action_then_exhausts() -> None:
    """Playback walks every record in order, then raises once past the end."""
    spec = _duel_spec()
    recorder, _report = _record(spec)
    log = recorder.log
    policy = replay(log)

    for record in log.records:
        assert policy(record.observation) == (record.action,)

    with pytest.raises(ReplayExhaustedError):
        policy(log.records[-1].observation)


def test_replay_log_rejects_a_duplicate_call_seq() -> None:
    """Adding two records with the same ``call_seq`` is an ambiguity error."""
    spec = _duel_spec()
    recorder, _report = _record(spec)
    record = recorder.log.records[-1]

    with pytest.raises(ReplayError, match="duplicate call_seq"):
        recorder.log.add(
            ReplayRecord(
                call_seq=record.call_seq,
                observation=record.observation,
                action=record.action,
            )
        )


def test_replay_log_rejects_an_out_of_order_call_seq() -> None:
    """A ``call_seq`` below the last one is an ambiguity error (not a duplicate)."""
    spec = _duel_spec()
    recorder, _report = _record(spec)
    record = recorder.log.records[0]

    with pytest.raises(ReplayError, match="is not after"):
        recorder.log.add(
            ReplayRecord(
                call_seq=record.call_seq,
                observation=record.observation,
                action=record.action,
            )
        )


def test_from_jsonl_rejects_garbage() -> None:
    """An empty or unparseable payload raises ``ReplayError``, not a raw error."""
    with pytest.raises(ReplayError):
        ReplayLog.from_jsonl("")
    with pytest.raises(ReplayError):
        ReplayLog.from_jsonl("{not json}")