#!/usr/bin/env python3
"""
test_run.py -- regression tests for the decision log in run.py.

The log is the only record of what the model was asked and what it answered, so
two things must hold or every later analysis is wrong:

  * a model decision's row carries that decision's answers;
  * a decision that asked NOTHING carries none -- no inheritance from the
    previous step, which is what `client.last = None` before each decide() is
    for, and which fails silently if it is ever dropped.

Runs offline: a stub client stands in for JevClient and a fake sts2 module
stands in for the game, so no network and no running game are needed.

Run: python3 test_run.py
"""

from __future__ import annotations

import json
import pathlib
import sys
import tempfile
import types

import brain
import jev
import run
from jev import ChoiceAnswer, JevResponse, NoulAnswer, ScoreAnswer

PASS = 0
FAIL = 0


def check(label: str, got, want) -> None:
    global PASS, FAIL
    if got == want:
        PASS += 1
        print(f"  ok   {label}: {got!r}")
    else:
        FAIL += 1
        print(f"  FAIL {label}: got {got!r}, want {want!r}")


class StubClient:
    """Answers every question locally, so the log can be tested without a model."""

    model = "stub"

    def __init__(self, noul: float = 0.9):
        self.noul = noul
        self.calls = 0

    def __repr__(self) -> str:
        return "StubClient()"

    def ask(self, state, questions, model=None):
        self.calls += 1
        answers = {}
        for qid, q in questions.items():
            if q.get("type") == "choice":
                options = list((q.get("criteria") or {}).keys())
                answers[qid] = ChoiceAnswer(
                    choice=options[0] if options else "",
                    probabilities={o: (0.9 if o == options[0] else 0.05)
                                   for o in options},
                    confidence=0.9,
                )
            elif q.get("type") == "noul":
                answers[qid] = NoulAnswer(noul=self.noul)
            else:
                answers[qid] = ScoreAnswer(score=1.0, confidence=0.9)
        return JevResponse(answers=answers, model="stub", input_tokens=0,
                           output_tokens=0, latency_s=0.0)


class FakeSts2:
    """A fixed sequence of states, then the last one forever."""

    class Sts2Error(RuntimeError):
        pass

    def __init__(self, states):
        self.states = states
        self.i = 0

    def is_up(self) -> bool:
        return True

    def state(self) -> dict:
        state = self.states[min(self.i, len(self.states) - 1)]
        self.i += 1
        return state

    def act(self, *a, **k):
        return types.SimpleNamespace(ok=True, message="")


print("=== 1. answer_record shapes (what the log stores) ===")
check("noul", jev.answer_record(NoulAnswer(0.72)),
      {"kind": "noul", "noul": 0.72, "yes": True, "gated": True})
check("noul gate outcome is recorded", jev.answer_record(NoulAnswer(0.52))["gated"], False)
choice = jev.answer_record(ChoiceAnswer("a", {"a": 0.7, "b": 0.2}, 0.5))
check("choice margin", choice["margin"], 0.5)
check("choice gated", choice["gated"], True)
check("choice probabilities survive", choice["probabilities"], {"a": 0.7, "b": 0.2})
check("score", jev.answer_record(ScoreAnswer(1.4, {}, {}, 0.9))["kind"], "score")

print()
print("=== 2. a decision that asked nothing logs jev: null ===")
stub = run.RecordingClient(StubClient())
code_decision = brain.Decision("end_turn", {}, "no playable cards", "code")
row = run.decision_record(1, "monster", {"act": 1, "floor": 1}, code_decision,
                          None, stub)
check("jev is null", row["jev"], None)
check("the session id travels with the row", row["session"], run.SESSION_ID)
check("the action is still recorded", row["action"], "end_turn")

print()
print("=== 3. end to end: a model row keeps its answers, the next row does not ===")
# Step 1 is a card_reward, which always asks the model. Step 2 is the main menu,
# which never does. If the reset is dropped, row 2 inherits row 1's answers --
# the failure this test exists to catch.
card_reward = {
    "state_type": "card_reward",
    "card_reward": {"can_skip": True, "cards": [
        {"id": "A", "name": "Stomp", "type": "Attack", "cost": "3",
         "description": "Deal 12 damage.", "rarity": "Uncommon",
         "is_upgraded": False, "index": 0},
        {"id": "B", "name": "Inflame", "type": "Power", "cost": "1",
         "description": "Gain 2 Strength.", "rarity": "Uncommon",
         "is_upgraded": False, "index": 1},
    ]},
    "run": {"act": 1, "floor": 2, "ascension": 0},
    "player": {"character": "The Ironclad", "hp": 70, "max_hp": 80, "gold": 99},
}
menu = {"state_type": "menu", "menu_screen": "main",
        "options": ["singleplayer", "quit"], "run": None}

tmp = pathlib.Path(tempfile.mkdtemp())
real_sts2, real_client = run.sts2, run.JevClient
# preflight() reads one state before the loop starts, so the sequence leads with
# a menu: preflight consumes that, step 1 sees the card_reward.
run.sts2 = FakeSts2([menu, card_reward, menu, menu])
run.JevClient = StubClient
argv = sys.argv
sys.argv = ["run.py", "--steps", "2", "--dry-run", "--capture-dir", str(tmp)]
try:
    rc = run.main()
finally:
    run.sts2, run.JevClient, sys.argv = real_sts2, real_client, argv

check("main() completed", rc, 0)
rows = [json.loads(line) for line in (tmp / "decisions.jsonl").read_text().splitlines()]
check("one row per decided action", len(rows), 2)
check("row 1 came from the model", rows[0]["source"], "jev")
check("...and carries its answers", sorted(rows[0]["jev"]["answers"]),
      ["good_card0", "good_card1", "skip_all"])
check("...with its value and gate outcome, not just a yes/no",
      (rows[0]["jev"]["answers"]["good_card0"]["noul"],
       rows[0]["jev"]["answers"]["good_card0"]["gated"]), (0.9, True))
check("row 2 asked nothing", rows[1]["jev"], None)
check("...and is a code decision", rows[1]["source"], "code")
check("every row carries the session", {r["session"] for r in rows}, {run.SESSION_ID})

print()
print("=== 4. the session row maps to an outcome (join half) ===")
# Real run files if this machine has them; skipped otherwise.
if run.HISTORY_DIR.exists():
    names = sorted(run.history_snapshot())
    check("history records visible", len(names) > 0, True)
    if names:
        outcome = run.run_outcome(names[-1])
        check("outcome names the file", outcome["file"], names[-1])
        check("outcome has the killer", "killed_by" in outcome, True)
        check("outcome has the deck size", isinstance(outcome["deck_size"], int), True)
else:
    print("  (no history directory on this machine, skipped)")

print()
print(f"=== {PASS} passed, {FAIL} failed ===")
sys.exit(1 if FAIL else 0)