from proteus.game.agents import VanillaAgent
from proteus.providers import FakeProvider
from proteus.game.runtime import SessionRunner, SessionTrace


def _run(response):
    agent = VanillaAgent(FakeProvider(responses=[response], model_name="fake-1"))
    return SessionRunner(
        "template", agent, seed=42, play_turns=3, use_probe=False,
    ).run()


def test_act_tokens_and_raw_text_persisted_per_turn():
    # FakeProvider sets output_tokens = word count; inline <think> gives thinking_tokens.
    trace = _run("<think>predator is two cells east</think>ACTION: up")
    t0 = trace.turns[0]
    assert t0.raw_text == "<think>predator is two cells east</think>ACTION: up"
    assert t0.output_tokens > 0                     # propagated from CompletionResult
    assert t0.input_tokens == 0                     # FakeProvider reports 0; output_tokens below is the live propagation check
    assert t0.thinking_tokens == 5                  # whitespace-split of the think block
    # use_probe=False: probe fields stay at their defaults.
    assert t0.probe_raw_text == ""
    assert t0.probe_thinking_tokens == 0


def test_accounting_survives_jsonl_roundtrip():
    trace = _run("<think>a b c</think>ACTION: up")
    reloaded = SessionTrace.model_validate_json(trace.model_dump_json())
    assert reloaded.model_dump() == trace.model_dump()
    assert reloaded.turns[0].thinking_tokens == 3
    assert reloaded.turns[0].raw_text == "<think>a b c</think>ACTION: up"


def test_probe_accounting_persisted_when_enabled():
    # Distinct probe vs act responses (different <think> word counts) so a
    # probe/act field cross-wiring bug would be visible: turn 1's probe consumes
    # responses[0] (3-word think) and turn 1's act consumes responses[1] (5-word think).
    agent = VanillaAgent(FakeProvider(
        responses=[
            "<think>it is east</think>predator is east",          # probe: 3 think-words
            "<think>move up and away now</think>I should go up\nACTION: up",  # act: 5 think-words
        ],
        model_name="fake-1",
    ))
    trace = SessionRunner(
        "template", agent, seed=42, play_turns=3, use_probe=True,
    ).run()
    t0 = trace.turns[0]
    assert t0.probe_a                               # answer recorded
    assert "it is east" in t0.probe_reasoning       # probe reasoning preserved
    assert t0.probe_raw_text                        # raw probe output preserved
    assert t0.probe_thinking_tokens == 3            # from the PROBE response
    assert t0.thinking_tokens == 5                  # from the ACT response — NOT cross-wired
    assert t0.probe_output_tokens > 0
    # survives JSONL round-trip
    reloaded = SessionTrace.model_validate_json(trace.model_dump_json())
    assert reloaded.model_dump() == trace.model_dump()