#!/usr/bin/env python3
"""
bakeoff.py — model JSON-reliability bake-off for the First Contact game.

Purpose
-------
Day-2 of the build is gated on one empirical question: which <=32B instruct model
emits clean, schema-valid JSON for the §3 prompt *reliably*? This harness answers
that in one run per candidate, instead of eyeballing outputs.

What it measures, in priority order (matches SPEC §3 mitigations + §11 Day-2):
  1. valid-JSON-on-FIRST-TRY rate   <-- the headline. NO retry. See note below.
  2. schema/action validity rate    (verb in allowed set, args well-formed)
  3. action presence rate           (did it produce a usable action at all)
  4. latency  (mean / median / p90) <-- feeds @spaces.GPU(duration=) + quota math

CRITICAL — why this bypasses the retry path
--------------------------------------------
parsing.py implements the §4 "retry-once -> safe wait" path. That path exists to
keep the *game* robust, and it works by silently repairing bad model output. If
this harness used the full engine turn loop, the retry would mask the exact
failures we want to count, and every model would look great. So we call the brain's
`respond()` to get RAW text and run a parse-WITHOUT-retry on it. `parse_once` below
is wired to the INNER extract+validate functions of parsing.py, never the
engine-facing `converse()` wrapper.

Sampling note
-------------
LocalBrain.respond samples (temperature=0.9 by default, per the 2026-06 sweep) —
so the SAME prompt yields different text each call, and JSON reliability is a
*rate*, not a yes/no. Use --repeats N (>=3 recommended on real models) to draw
each battery prompt several times.

Run modes (same stub/local discipline as the game loop)
-------------------------------------------------------
  --self-test            Score a built-in fixture brain (good/bad/wrapped/broken
                         JSON), assert the scorer classifies them correctly. ZERO
                         GPU, no model. Run FIRST to prove harness logic.
  --make-battery PATH    Generate battery.jsonl from the real headless arc (ZERO
                         GPU, StubBrain) so prompts span the true game distribution
                         (empty ledger -> populated -> generalization states).
  --brain stub           Replay the battery through StubBrain end to end (ZERO GPU
                         dry run of load->respond->score->table).
  --models A,B --brain local
                         The real measurement: build each LocalBrain, replay the
                         battery, score + time. Run on the Space.
  --arc                  Also play the full 5-challenge arc per model (secondary
                         signal: does it actually win, and at what retry/fallback
                         cost). Works on stub too.
  --arc-transcript       --arc plus a readable per-challenge transcript: your
                         utterance -> action chosen -> won? -> the alien's reply,
                         gap and proposed concept, with the generalization beats
                         called out. Each challenge gets up to ARC_MAX_TURNS (3)
                         turns with a neutral "Go on." nudge — real play is
                         multi-turn. THE check for "do surprise/secret actually
                         land, and is the voice worth shipping".
  --temps 0.0,0.3,0.5,0.7,1.0
                         THE decision tool. One model load; per temperature, run
                         the battery + an arc pass and report reliability AND two
                         voice-liveliness proxies (distinct-utterance ratio,
                         unique candidate_concept count) AND arc-win. The table
                         shows JSON reliability and voice liveliness moving in
                         opposite directions so the knee is visible — answering
                         "one temp for both jobs, or decouple via constrained
                         decoding?". Set LocalBrain temp via LOCALBRAIN_TEMPERATURE
                         (default 0.9) for normal runs.

Battery line schema: {"prompt": "...", "challenge_id": "...", "turn": N}
(only "prompt" is required; the rest enable per-challenge grouping.)

Stdlib only. torch is NOT imported unless a real LocalBrain is actually built.
"""

from __future__ import annotations

import argparse
import json
import statistics
import sys
import time
from dataclasses import dataclass, field
from typing import Optional

# Real parser internals — light, torch-free. parse_once wires to THESE so the
# rates reflect exactly what the game's parser/validator will accept.
from game.parsing import ParseError, extract_json, validate_action
from game.world import initial_world

# A reference world for action referential checks. Object/agent/container ids are
# constant across every challenge (the world resets to initial_world each time),
# so validating any battery action against this one world is correct.
_REF_WORLD = initial_world()


# =============================================================================
# ADAPTER BLOCK  —  reconciled against game/brain.py + game/parsing.py (Day 1).
# =============================================================================

# --- 1) Building a brain -----------------------------------------------------
# game/brain.py exposes StubBrain() and LocalBrain(model_id=...); the Brain method
# is `.respond(prompt) -> str` returning RAW model text. (LocalBrain also reads
# MODEL_ID from env when model_id is None, but here we pass it explicitly.)
def make_brain(kind: str, model_id: Optional[str] = None):
    """kind: 'local' | 'stub'. Returns an object with .respond(prompt)->str."""
    if kind == "stub":
        from game.brain import StubBrain
        return StubBrain()
    if kind == "local":
        from game.brain import LocalBrain          # imports torch lazily, on build
        return LocalBrain(model_id=model_id)
    if kind == "modal":
        from game.brain import ModalBrain
        return ModalBrain(model_id=model_id)   # endpoint from MODAL_ENDPOINT
    raise ValueError(f"unknown brain kind: {kind}")


# --- 2) Parse WITHOUT retry --------------------------------------------------
# Wired to the real tolerant extractor + validator. extract_json(raw) returns the
# first balanced {...} block as a STRING (or None); validate_action(action, world)
# raises ParseError on an unknown verb / missing args / bad reference. Neither
# touches converse()'s retry/fallback, which is the whole point.
@dataclass
class ParseOutcome:
    json_ok: Optional[bool]     # did a JSON object parse out of the raw text?
    schema_ok: Optional[bool]   # required fields present + types right?
    action_ok: Optional[bool]   # verb in allowed set AND args valid for that verb?
    raw_len: int = 0            # chars of raw output (sanity / verbosity signal)
    # The two voice-liveliness surfaces (SPEC §3): the prose and the proposed
    # concept. Captured so the --temps sweep can measure diversity, not just JSON.
    utterance: Optional[str] = None
    candidate_key: Optional[str] = None   # normalized candidate_concept fingerprint


def parse_once(raw_text: str) -> ParseOutcome:
    """Single-shot parse+validate. NO RETRY. Mirrors parse_response's field rules
    (§3) but reports each stage separately instead of raising on the first miss."""
    raw_len = len(raw_text)

    blob = extract_json(raw_text)
    if blob is None:
        return ParseOutcome(json_ok=False, schema_ok=None, action_ok=None, raw_len=raw_len)
    try:
        obj = json.loads(blob)
    except json.JSONDecodeError:
        return ParseOutcome(json_ok=False, schema_ok=None, action_ok=None, raw_len=raw_len)
    if not isinstance(obj, dict):
        return ParseOutcome(json_ok=False, schema_ok=None, action_ok=None, raw_len=raw_len)

    # Schema: same required fields/types parse_response enforces (minus the world).
    has_action = isinstance(obj.get("action"), dict)
    utt = obj.get("utterance")
    has_utterance = isinstance(utt, str) and bool(utt.strip())
    gap = obj.get("gap")
    gap_ok = gap is None or isinstance(gap, str)
    cand = obj.get("candidate_concept")
    cand_ok = cand is None or isinstance(cand, dict)
    schema_ok = bool(has_action and has_utterance and gap_ok and cand_ok)

    # Action validity through the REAL validator (canonical verb set + arg/ref checks).
    if not has_action:
        action_ok = False
    else:
        try:
            validate_action(obj["action"], _REF_WORLD)
            action_ok = True
        except ParseError:
            action_ok = False

    # Liveliness surfaces: the utterance text, and a normalized fingerprint of the
    # proposed concept (its understanding gloss — the inventive bit).
    utterance_out = utt.strip() if (isinstance(utt, str) and utt.strip()) else None
    candidate_key = None
    if isinstance(cand, dict):
        raw_key = cand.get("understanding") or cand.get("label") or cand.get("id") or ""
        norm = " ".join(str(raw_key).split()).lower()
        candidate_key = norm or None

    return ParseOutcome(
        json_ok=True, schema_ok=schema_ok, action_ok=action_ok, raw_len=raw_len,
        utterance=utterance_out, candidate_key=candidate_key,
    )


# --- 3) Secondary metric: does the full arc still win, at what retry cost? ----
# Wired to game.engine (the same loop tests/test_loop_stub.py exercises). Retries
# are counted without touching converse(): a _CountingBrain tallies respond()
# calls; normal turns call respond once, a §4 retry makes it twice, so
# retries = respond_calls - turns_played. Fallbacks come from used_fallback.
_ARC_UTTERANCES = {
    "warmup": "Put the red stone in the basket.",
    "hide": "Hide the blue stone from the other one.",
    "gift": "Give the other one a present.",
    "surprise": "Make a surprise for the other one.",
    "secret": "Keep your blue stone a secret from the other one.",
}


class _CountingBrain:
    def __init__(self, inner):
        self.inner = inner
        self.calls = 0

    def respond(self, prompt: str) -> str:
        self.calls += 1
        return self.inner.respond(prompt)


# A real player keeps talking until the goal is met, so the arc runner gets up to
# ARC_MAX_TURNS per challenge, nudging with a neutral follow-up that teaches
# nothing. (The old one-turn-per-challenge scored multi-step plans — pick up,
# THEN give — as failures, even though the actual game allows them: the 2026-06
# transcript showed 14B "losing" gift/surprise exactly this way.)
ARC_MAX_TURNS = 3
_ARC_FOLLOWUP = "Go on."


def play_arc(brain) -> Optional["ArcStats"]:
    """Play all 5 challenges with `brain`, up to ARC_MAX_TURNS turns each (real
    play is multi-turn). Reports wins + retry/fallback cost AND a per-challenge
    transcript (every turn: what the alien said/did/proposed) — a scalar like
    3/5 can hide that BOTH generalization beats failed, which is the go/no-go
    question for the demo."""
    from game.challenges import CHALLENGES
    from game.engine import advance_challenge, confirm_candidate, new_session, run_turn

    counter = _CountingBrain(brain)
    session = new_session()
    transcript: list[dict] = []
    fallbacks = turns_played = 0
    for i, ch in enumerate(CHALLENGES):
        entry = {
            "challenge_id": ch.id,
            "teaches": ch.teaches,
            "relies_on": list(ch.relies_on),
            "won": False,
            "turns_used": 0,
            "turns": [],
        }
        for turn_no in range(1, ARC_MAX_TURNS + 1):
            utterance = _ARC_UTTERANCES.get(ch.id, "wait") if turn_no == 1 else _ARC_FOLLOWUP
            calls_before = counter.calls
            res = run_turn(session, utterance, counter)
            turns_played += 1
            fell_back = bool(getattr(res.response, "used_fallback", False))
            if fell_back:
                fallbacks += 1
            cand = res.response.candidate_concept
            confirmed = False
            if res.learn_offer:
                # Mirror a player pressing "Yes, it learned that" — later turns
                # then see the concept in the ledger, exactly like real play.
                confirmed = confirm_candidate(session) is not None
            entry["turns"].append({
                "player": utterance,
                "action": {"verb": res.response.action.verb,
                           "args": dict(res.response.action.args)},
                "alien": res.response.utterance,
                "gap": res.response.gap,
                "candidate_concept": cand if isinstance(cand, dict) else None,
                "confirmed": confirmed,
                "reapplied": list(res.reapplied),
                "retries": max(0, counter.calls - calls_before - 1),
                "fallback": fell_back,
            })
            entry["turns_used"] = turn_no
            if res.won:
                entry["won"] = True
                break
        transcript.append(entry)
        if i < len(CHALLENGES) - 1:
            advance_challenge(session)

    won_ids = [e["challenge_id"] for e in transcript if e["won"]]
    gen_entries = [e for e in transcript if e["teaches"] is None and e["relies_on"]]
    return ArcStats(
        challenges_won=len(won_ids),
        challenges_total=len(CHALLENGES),
        retries_triggered=max(0, counter.calls - turns_played),
        fallbacks_hit=fallbacks,
        won_ids=won_ids,
        gen_won=sum(1 for e in gen_entries if e["won"]),
        gen_total=len(gen_entries),
        transcript=transcript,
    )


@dataclass
class ArcStats:
    challenges_won: int
    challenges_total: int
    retries_triggered: int
    fallbacks_hit: int
    won_ids: list = field(default_factory=list)
    gen_won: int = 0     # generalization challenges won — the demo moments
    gen_total: int = 0
    transcript: list = field(default_factory=list)   # per-challenge detail dicts

# =============================================================================
# END ADAPTER BLOCK
# =============================================================================


@dataclass
class ModelReport:
    model_id: str
    n: int = 0
    json_ok: int = 0
    schema_ok: int = 0
    action_ok: int = 0
    schema_na: int = 0          # parser couldn't judge this stage
    action_na: int = 0
    latencies: list[float] = field(default_factory=list)
    errors: int = 0             # respond() threw
    per_challenge: dict = field(default_factory=dict)  # challenge_id -> [json_ok bools]
    arc: Optional[ArcStats] = None
    temperature: Optional[float] = None                # set in --temps sweep rows
    utt_by_prompt: dict = field(default_factory=dict)  # prompt -> [utterances]
    candidate_keys: list = field(default_factory=list)  # all candidate fingerprints

    def _rate(self, num: int, denom: int) -> str:
        return f"{(100.0 * num / denom):5.1f}%" if denom else "  n/a"

    def distinct_utterance_ratio(self) -> float:
        """Mean over prompts of unique/total utterances across the repeats. HIGH =
        lively voice; LOW = the alien saying the same dutiful line every draw.
        Moves opposite to JSON reliability as temperature drops. (Needs repeats>1
        to be informative — at repeats=1 every prompt is trivially 1.0.)"""
        ratios = [len(set(u)) / len(u) for u in self.utt_by_prompt.values() if u]
        return statistics.mean(ratios) if ratios else 0.0

    def unique_candidates(self) -> int:
        """Distinct candidate_concept proposals seen — the inventive surface."""
        return len(set(self.candidate_keys))

    def summary_row(self) -> list[str]:
        lat = self.latencies
        med = statistics.median(lat) if lat else 0.0
        p90 = (sorted(lat)[max(0, int(len(lat) * 0.9) - 1)] if lat else 0.0)
        return [
            self.model_id,
            str(self.n),
            self._rate(self.json_ok, self.n),
            self._rate(self.schema_ok, self.n - self.schema_na),
            self._rate(self.action_ok, self.n - self.action_na),
            f"{med:5.2f}s",
            f"{p90:5.2f}s",
            str(self.errors),
        ]


def score_brain_over_battery(brain, battery: list[dict], label: str,
                             warmup: int = 1, repeats: int = 1,
                             verbose: bool = False) -> ModelReport:
    """Replay every prompt in the battery (x `repeats`) through `brain`, scoring
    first-try parse. The first `warmup` calls are discarded from latency stats."""
    rep = ModelReport(model_id=label)

    # Warmup: the first ZeroGPU call pays cold-start + GPU allocation; exclude it
    # so the numbers reflect steady state.
    for i in range(min(warmup, len(battery))):
        try:
            brain.respond(battery[i]["prompt"])
        except Exception:
            pass

    for item in battery:
        prompt = item["prompt"]
        cid = item.get("challenge_id", "_")
        for _ in range(max(1, repeats)):
            t0 = time.perf_counter()
            try:
                raw = brain.respond(prompt)
            except Exception as e:
                rep.errors += 1
                if verbose:
                    print(f"  [respond error] {e}", file=sys.stderr)
                continue
            rep.latencies.append(time.perf_counter() - t0)
            rep.n += 1

            out = parse_once(raw)
            if out.json_ok:
                rep.json_ok += 1
            if out.schema_ok is None:
                rep.schema_na += 1
            elif out.schema_ok:
                rep.schema_ok += 1
            if out.action_ok is None:
                rep.action_na += 1
            elif out.action_ok:
                rep.action_ok += 1

            rep.per_challenge.setdefault(cid, []).append(bool(out.json_ok))
            if out.utterance:
                rep.utt_by_prompt.setdefault(prompt, []).append(out.utterance)
            if out.candidate_key:
                rep.candidate_keys.append(out.candidate_key)

            if verbose and not out.json_ok:
                snippet = raw[:160].replace("\n", " ")
                print(f"  [first-try FAIL @ {cid}] {snippet!r}", file=sys.stderr)

    return rep


def _print_grid(headers: list[str], rows: list[list[str]]) -> None:
    widths = [max(len(headers[i]), *(len(r[i]) for r in rows)) if rows
              else len(headers[i]) for i in range(len(headers))]

    def fmt(cells):
        return "  ".join(c.ljust(widths[i]) for i, c in enumerate(cells))

    print(fmt(headers))
    print("  ".join("-" * w for w in widths))
    for row in rows:
        print(fmt(row))


def print_table(reports: list[ModelReport]) -> None:
    headers = ["MODEL", "N", "JSON-1st", "SCHEMA", "ACTION", "MED", "P90", "ERR"]
    print()
    _print_grid(headers, [r.summary_row() for r in reports])
    print()
    print("JSON-1st = valid JSON on first try, NO retry  <-- rank on this first.")
    print("SCHEMA/ACTION computed over prompts the parser could judge (n/a excluded).")


def print_per_challenge(reports: list[ModelReport]) -> None:
    """Where does JSON reliability break down? Often the richer (later) prompts."""
    for r in reports:
        if not r.per_challenge:
            continue
        print(f"\nJSON-1st by challenge - {r.model_id}:")
        for cid, oks in r.per_challenge.items():
            rate = 100.0 * sum(oks) / len(oks) if oks else 0.0
            print(f"  {cid:12s} {rate:5.1f}%  ({sum(oks)}/{len(oks)})")


def print_arc(reports: list[ModelReport]) -> None:
    any_arc = any(r.arc for r in reports)
    if not any_arc:
        return
    print("\nArc play-through (secondary):")
    for r in reports:
        a = r.arc
        if not a:
            continue
        lost = [e["challenge_id"] for e in a.transcript if not e["won"]]
        print(f"  {r.model_id:32s}  won {a.challenges_won}/{a.challenges_total}  "
              f"gen {a.gen_won}/{a.gen_total}  retries={a.retries_triggered}  "
              f"fallbacks={a.fallbacks_hit}  lost: {','.join(lost) or '-'}")


def _safe(s: object) -> str:
    """Console-safe text: model output is unicode, but Windows pipes are often
    cp1252 — replace what the active stdout encoding can't represent instead of
    letting print() raise mid-report."""
    enc = getattr(sys.stdout, "encoding", None) or "utf-8"
    return str(s).encode(enc, "replace").decode(enc)


def print_arc_transcript(rep: ModelReport) -> None:
    """The eyeball check the scalars can't give: per challenge, what the alien
    actually did and said — and whether the generalization beats landed."""
    a = rep.arc
    if not a or not a.transcript:
        return
    print(f"\nArc transcript - {rep.model_id}")
    total = len(a.transcript)
    for i, e in enumerate(a.transcript, 1):
        if e["teaches"]:
            kind = f"teaches {e['teaches']}"
        elif e["relies_on"]:
            kind = "GENERALIZATION via " + "+".join(e["relies_on"])
        else:
            kind = "mechanical"
        status = (f"WON in {e['turns_used']}" if e["won"]
                  else f"LOST after {e['turns_used']} <--")
        print(f"\n[{i}/{total}] {e['challenge_id']} ({kind})  {status}")
        for t_no, t in enumerate(e["turns"], 1):
            print(f"  T{t_no} you    > {_safe(t['player'])}")
            args = ", ".join(f"{k}={v}" for k, v in t["action"]["args"].items())
            print(f"     action > {t['action']['verb']}({args})")
            print(f"     alien  > {_safe(t['alien'])}")
            if t["gap"]:
                print(f"     gap    > {_safe(t['gap'])}")
            cand = t["candidate_concept"]
            if cand:
                tag = "confirmed -> ledger" if t["confirmed"] else "not added (duplicate/rejected)"
                print(f"     learn  > {cand.get('id', '?')}: "
                      f"\"{_safe(cand.get('understanding', ''))}\"  [{tag}]")
            extras = []
            if t["reapplied"]:
                extras.append("reapplied: " + ",".join(t["reapplied"]))
            if t["retries"]:
                extras.append(f"retries={t['retries']}")
            if t["fallback"]:
                extras.append("FALLBACK")
            if extras:
                print(f"     note   > {'  '.join(extras)}")


def print_quota_math(reports: list[ModelReport], daily_seconds: int = 2400) -> None:
    """40 min/day org quota = 2400s. Translate latency into turns/day + duration=."""
    print("\nQuota & duration guidance (org ZeroGPU = 40 min/day = 2400s effective):")
    for r in reports:
        if not r.latencies:
            continue
        med = statistics.median(r.latencies)
        p90 = sorted(r.latencies)[max(0, int(len(r.latencies) * 0.9) - 1)]
        if med < 1e-3:   # stub / instant — latency-based guidance is meaningless
            print(f"  {r.model_id:32s}  (instant — no GPU; run with --brain local)")
            continue
        turns = int(daily_seconds / med)
        # duration= should cover the slow tail with margin but stay tight for queue
        # priority; p90 + ~25%, rounded up, is a sane start.
        suggested = int(p90 * 1.25) + 1
        print(f"  {r.model_id:32s}  ~{turns:5d} turns/day  "
              f"suggested @spaces.GPU(duration={suggested})")
    print()


# ----------------------------------------------------------------------------
# Temperature sweep — the decision tool (SPEC §3 mitigations, §11 Day-2)
# ----------------------------------------------------------------------------
# The sweep answers "is there ONE temperature where JSON is reliable enough AND
# the voice still has life?" — not "what temp gives the best JSON". So it reports
# reliability AND two liveliness proxies AND arc-win, per temperature, in one
# table, so the tradeoff (and its knee) is visible rather than inferred.
def _med_p90(lat: list[float]) -> tuple[float, float]:
    if not lat:
        return 0.0, 0.0
    return statistics.median(lat), sorted(lat)[max(0, int(len(lat) * 0.9) - 1)]


def run_sweep(brain, model_id: str, battery: list[dict], temps: list[float],
              repeats: int, warmup: int, verbose: bool) -> list[ModelReport]:
    """Reuse ONE loaded model; per temperature run the formatting battery AND an
    arc pass. The sampler is retuned via brain.temperature, never reloaded."""
    # Preflight: one COLD call (pays the one-time model->GPU attach), then one
    # WARM call whose latency drives the estimate. Estimating from the cold call
    # overstated cost ~40x in practice (it once cried "1160% of budget" for a run
    # that actually cost ~12%) — a scary-wrong number is worse than no number.
    cold = warm = None
    if warmup > 0 and battery:
        for phase in ("cold", "warm"):
            t0 = time.perf_counter()
            try:
                brain.respond(battery[0]["prompt"])
            except Exception as e:
                if verbose:
                    print(f"  [{phase} warmup error] {e}", file=sys.stderr)
                break
            if phase == "cold":
                cold = time.perf_counter() - t0
            else:
                warm = time.perf_counter() - t0

    # Upper bound: challenges that win early use fewer than ARC_MAX_TURNS.
    arc_calls = len(_ARC_UTTERANCES) * ARC_MAX_TURNS
    total = len(temps) * (len(battery) * repeats + arc_calls)
    if warm is not None:
        est = total * warm
        print(f"[preflight] cold first call ~{cold:.1f}s (one-time load), "
              f"warm ~{warm:.2f}s/call -> ~{total} calls x {warm:.2f}s = ~{est:.0f}s "
              f"(~{100.0 * est / 2400:.0f}% of the 2400s/day budget; retries add a "
              f"little). Ctrl-C now if that is too much.")
    else:
        print(f"[preflight] ~{total} GPU calls planned ({len(temps)} temps x "
              f"({len(battery)}x{repeats} battery + ~{arc_calls} arc)); "
              f"could not time a warmup call.")

    reports = []
    for t in temps:
        brain.temperature = t                   # retune sampler; model stays loaded
        rep = score_brain_over_battery(brain, battery, label=f"{model_id} @T={t:g}",
                                       warmup=0, repeats=repeats, verbose=verbose)
        rep.temperature = t
        rep.arc = play_arc(brain)               # arc at THIS temperature
        reports.append(rep)
    return reports


def print_sweep_table(reports: list[ModelReport], model_id: str) -> None:
    reports = sorted(reports, key=lambda r: (r.temperature if r.temperature is not None else 0.0))
    headers = ["TEMP", "N", "JSON-1st", "SCHEMA", "ACTION",
               "UTTER-DIV", "UNIQ-CC", "ARC", "GEN", "RETRY", "FALLBK", "MED", "P90"]
    rows = []
    for r in reports:
        med, p90 = _med_p90(r.latencies)
        a = r.arc
        rows.append([
            f"{r.temperature:g}" if r.temperature is not None else "?",
            str(r.n),
            r._rate(r.json_ok, r.n),
            r._rate(r.schema_ok, r.n - r.schema_na),
            r._rate(r.action_ok, r.n - r.action_na),
            f"{100.0 * r.distinct_utterance_ratio():5.1f}%",
            str(r.unique_candidates()),
            f"{a.challenges_won}/{a.challenges_total}" if a else "n/a",
            f"{a.gen_won}/{a.gen_total}" if a else "n/a",
            str(a.retries_triggered) if a else "-",
            str(a.fallbacks_hit) if a else "-",
            f"{med:5.2f}s",
            f"{p90:5.2f}s",
        ])
    print(f"\nTemperature sweep - {model_id}\n")
    _print_grid(headers, rows)
    print()
    print("JSON-1st climbs as TEMP falls; UTTER-DIV (voice) + UNIQ-CC (invention) fall with it.")
    print("UTTER-DIV = mean distinct-utterance ratio across repeats; UNIQ-CC = distinct proposals.")
    print("Find the KNEE: JSON-1st plateaus while liveliness is still dropping -> that TEMP wins.")
    print("ARC = arc wins at that TEMP; GEN = the generalization beats won (surprise/secret) --")
    print("the demo moments. Weight GEN most; --arc-transcript shows them play out in full.")
    print("If no TEMP gives reliable JSON AND a live voice -> decouple: constrained-decode the")
    print("envelope, keep TEMP warm. The sweep has then justified that dependency with data.")


def report_to_dict(r: ModelReport) -> dict:
    """Compact, machine-readable summary for the writeup. Omits the raw battery
    utterances / full-prompt keys (huge); the 5-entry arc transcript IS kept —
    it's the demo-quality evidence."""
    med, p90 = _med_p90(r.latencies)
    judged_schema = r.n - r.schema_na
    judged_action = r.n - r.action_na
    return {
        "model_id": r.model_id,
        "temperature": r.temperature,
        "n": r.n,
        "json_first_try_pct": round(100.0 * r.json_ok / r.n, 2) if r.n else None,
        "schema_pct": round(100.0 * r.schema_ok / judged_schema, 2) if judged_schema > 0 else None,
        "action_pct": round(100.0 * r.action_ok / judged_action, 2) if judged_action > 0 else None,
        "utterance_diversity": round(r.distinct_utterance_ratio(), 4),
        "unique_candidates": r.unique_candidates(),
        "median_latency_s": round(med, 3),
        "p90_latency_s": round(p90, 3),
        "errors": r.errors,
        "per_challenge_json_pct": {
            k: round(100.0 * sum(v) / len(v), 1) for k, v in r.per_challenge.items() if v
        },
        "arc": r.arc.__dict__ if r.arc else None,
    }


# ----------------------------------------------------------------------------
# Battery generation — from the real headless arc (ZERO GPU)
# ----------------------------------------------------------------------------
# Multiple natural phrasings per challenge widen the distribution; the canonical
# (first) phrasing is used to STEP the engine so state advances realistically.
_BATTERY_PHRASINGS = {
    "warmup": [
        "Put the red stone in the basket.",
        "Place the red stone into the basket.",
    ],
    "hide": [
        "Hide the blue stone from the other one.",
        "Keep the blue stone where the other one can't see it.",
        "Conceal the blue stone in the basket.",
    ],
    "gift": [
        "Give the other one a present.",
        "Hand the other one a gift.",
    ],
    "surprise": [
        "Make a surprise for the other one.",
        "Prepare a surprise for them.",
    ],
    "secret": [
        "Keep your blue stone a secret from the other one.",
        "Don't let the other one know about the blue stone — it's a secret.",
    ],
}


def make_battery(out_path: str) -> int:
    """Replay the arc with StubBrain (zero GPU), dumping each build_prompt() output
    at its real state. Returns the number of prompts written."""
    from game.brain import StubBrain
    from game.challenges import CHALLENGES
    from game.engine import advance_challenge, confirm_candidate, new_session, run_turn
    from game.prompt import build_prompt

    brain = StubBrain()
    session = new_session()
    rows: list[dict] = []
    for i, ch in enumerate(CHALLENGES):
        phrasings = _BATTERY_PHRASINGS.get(ch.id, [_ARC_UTTERANCES.get(ch.id, "")])
        for utter in phrasings:
            # build_prompt sees the CURRENT ledger/world (pre-step) — the real
            # situation as presented to the model that turn.
            prompt = build_prompt(session.ledger, session.world, ch, utter)
            rows.append({"prompt": prompt, "challenge_id": ch.id, "turn": session.turn})
        # Advance state with the canonical winning utterance, confirming concepts
        # so later challenges carry a populated ledger (the generalization states).
        res = run_turn(session, _ARC_UTTERANCES[ch.id], brain)
        if res.learn_offer:
            confirm_candidate(session)
        if i < len(CHALLENGES) - 1:
            advance_challenge(session)

    with open(out_path, "w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r) + "\n")
    return len(rows)


# ----------------------------------------------------------------------------
# Self-test: prove the SCORER is correct with zero GPU and no real model.
# ----------------------------------------------------------------------------
class _FixtureBrain:
    """Returns canned outputs by index: good, bad-verb, prose-wrapped, broken."""
    def __init__(self):
        self._canned = [
            # valid: parses, schema ok, action ok
            '{"action": {"verb": "put_in", "args": {"obj_id": "blue_stone", '
            '"container_id": "basket"}}, "utterance": "I place it inside.", '
            '"gap": null, "candidate_concept": null}',
            # valid JSON, valid schema, but BAD action verb -> action_ok False
            '{"action": {"verb": "teleport", "args": {}}, "utterance": "hm", '
            '"gap": null, "candidate_concept": null}',
            # prose wrapper around valid JSON -> tolerant extractor should still parse
            'Sure! Here is my move:\n{"action": {"verb": "wait", "args": {}}, '
            '"utterance": "..."}\nHope that helps!',
            # broken JSON -> json_ok False
            '{"action": {"verb": "give" , "utterance": OOPS no close',
        ]
        self._i = 0

    def respond(self, prompt: str) -> str:
        out = self._canned[self._i % len(self._canned)]
        self._i += 1
        return out


def run_self_test() -> int:
    battery = [{"prompt": "x", "challenge_id": "fixture"} for _ in range(4)]
    rep = score_brain_over_battery(_FixtureBrain(), battery, label="_fixture",
                                   warmup=0, repeats=1, verbose=False)
    # Over the 4 canned outputs: json_ok 3/4 (broken fails); action_ok 2
    # (good + wait ok; teleport rejected by the real validator).
    problems = []
    if rep.json_ok != 3:
        problems.append(f"expected json_ok=3, got {rep.json_ok}")
    if rep.n != 4:
        problems.append(f"expected n=4, got {rep.n}")
    if rep.action_ok != 2:
        problems.append(f"expected action_ok=2 (bad verb rejected), got {rep.action_ok}")

    print_table([rep])
    if problems:
        print("SELF-TEST FAILED:")
        for p in problems:
            print(f"  - {p}")
        print("\nThis means the scorer/adapter is miswired, NOT a model problem.")
        print("Fix parse_once() in the adapter block before measuring real models.")
        return 1
    print("SELF-TEST PASSED - scorer classifies good/bad/wrapped/broken correctly.")
    print("Safe to spend GPU on real measurement now.")
    return 0


def load_battery(path: str) -> list[dict]:
    battery = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            if "prompt" not in obj:
                raise ValueError(f"battery line missing 'prompt': {line[:80]}")
            battery.append(obj)
    if not battery:
        raise ValueError(f"battery file {path} is empty")
    return battery


def main() -> int:
    ap = argparse.ArgumentParser(description="Model JSON-reliability bake-off.")
    ap.add_argument("--self-test", action="store_true",
                    help="Validate the scorer with a fixture brain (zero GPU). Run first.")
    ap.add_argument("--make-battery", type=str, default="",
                    help="Generate a prompt battery JSONL from the arc (zero GPU) and exit.")
    ap.add_argument("--models", type=str, default="",
                    help="Comma-separated MODEL_IDs to measure via LocalBrain.")
    ap.add_argument("--battery", type=str, default="battery.jsonl",
                    help="Path to the prompt battery JSONL.")
    ap.add_argument("--repeats", type=int, default=1,
                    help="Draws per prompt (>=3 on real models; respond() samples).")
    ap.add_argument("--temps", type=str, default="",
                    help="Temperature sweep, e.g. 0.0,0.3,0.5,0.7,1.0. Reuses ONE model "
                         "load; per temp runs the battery + an arc pass and adds the "
                         "voice-liveliness columns. Spend repeats densely near the knee.")
    ap.add_argument("--warmup", type=int, default=1,
                    help="Calls to discard before timing (cold start / GPU alloc).")
    ap.add_argument("--arc", action="store_true",
                    help="Also play the full arc per model (wins + retry/fallback cost).")
    ap.add_argument("--arc-transcript", action="store_true",
                    help="--arc plus the readable per-challenge transcript (utterance -> "
                         "action -> won? -> reply/gap/proposal). The check for whether "
                         "the generalization beats actually land.")
    ap.add_argument("--verbose", action="store_true",
                    help="Print first-try failures to stderr for inspection.")
    ap.add_argument("--brain", type=str, default="local", choices=["local", "stub", "modal"],
                    help="Brain implementation to use (local, stub, modal).")
    args = ap.parse_args()
    if args.arc_transcript:
        args.arc = True

    if args.self_test:
        return run_self_test()

    if args.make_battery:
        n = make_battery(args.make_battery)
        print(f"Wrote {n} prompts to {args.make_battery}.")
        return 0

    battery = load_battery(args.battery)
    print(f"Loaded {len(battery)} prompts from {args.battery} "
          f"(x{args.repeats} repeats = {len(battery) * args.repeats} draws).")

    reports: list[ModelReport] = []
    sweep = bool(args.temps.strip())
    temps = [float(x) for x in args.temps.split(",") if x.strip() != ""] if sweep else []

    if args.brain == "stub":
        brain = make_brain("stub")
        if sweep:
            reports = run_sweep(brain, "StubBrain(dry-run)", battery, temps,
                                args.repeats, args.warmup, args.verbose)
            print_sweep_table(reports, "StubBrain(dry-run)")
            if args.arc_transcript:
                for r in reports:
                    print_arc_transcript(r)
        else:
            rep = score_brain_over_battery(brain, battery, label="StubBrain(dry-run)",
                                           warmup=args.warmup, repeats=args.repeats,
                                           verbose=args.verbose)
            if args.arc:
                rep.arc = play_arc(make_brain("stub"))
            reports = [rep]
    else:
        model_ids = [m.strip() for m in args.models.split(",") if m.strip()]
        if not model_ids:
            print("No --models given. Provide e.g. "
                  "--models Qwen/Qwen2.5-7B-Instruct,<other>", file=sys.stderr)
            return 2
        for mid in model_ids:
            print(f"\n=== Measuring {mid} ===")
            try:
                brain = make_brain(args.brain, model_id=mid)
            except Exception as e:
                print(f"  [build failed] {e}", file=sys.stderr)
                continue
            if sweep:
                rs = run_sweep(brain, mid, battery, temps,
                               args.repeats, args.warmup, args.verbose)
                print_sweep_table(rs, mid)
                if args.arc_transcript:
                    for r in rs:
                        print_arc_transcript(r)
                reports.extend(rs)
            else:
                rep = score_brain_over_battery(brain, battery, label=mid,
                                               warmup=args.warmup, repeats=args.repeats,
                                               verbose=args.verbose)
                if args.arc:
                    rep.arc = play_arc(brain)
                reports.append(rep)

    if not sweep:
        print_table(reports)
        print_per_challenge(reports)
        print_arc(reports)
        if args.arc_transcript:
            for r in reports:
                print_arc_transcript(r)
        print_quota_math(reports)

    with open("bakeoff_results.json", "w", encoding="utf-8") as f:
        json.dump([report_to_dict(r) for r in reports], f, indent=2)
    print("Wrote bakeoff_results.json")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())