#!/usr/bin/env python3
"""Capture real rollouts into the fixtures the site replays.

The site does no live inference. That is a deliberate architecture choice, not
a limitation: a public demo with no auth cannot hold an API key, a live call is
slow and flaky on conference wifi, and a recorded run can be scrubbed,
verified, permalinked and blind-compared in ways a live one cannot.

What makes it honest rather than a video is that the browser re-derives every
number it shows from the recorded moves, and says so. See verify.ts.

Usage:
  uv run python envs/capture.py --arm base-off --seeds 0-7
  uv run python envs/capture.py --arm solver   --seeds 0-7   # no model needed
  uv run python envs/capture.py --taskset alert-triage --arm base-on --seeds 1-7,9
  uv run python envs/capture.py --taskset alert-triage --arm targeted --seeds 1-7,9

spark-1 serves one model, single-stream: run the model arms one after another,
never concurrently. A thinking arm on alert-triage takes minutes per seed —
run it in the background with stdout redirected to a file and poll the
fixtures being written rather than the log.
"""

from __future__ import annotations

import argparse
import json
import math
import os
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent / "wordle_five"))
sys.path.insert(0, str(Path(__file__).parent / "alert_triage"))

import numpy as np  # noqa: E402

from alert_triage import policies as AT  # noqa: E402
from alert_triage import taskset as AT_taskset  # noqa: E402
from alert_triage.generator import is_held_out, world_for_seed  # noqa: E402

from wordle_five import solver as S  # noqa: E402
from wordle_five.engine import MAX_GUESSES, Game, answers  # noqa: E402
from wordle_five.protocol import (  # noqa: E402
    parse_guess,
    render_feedback,
    render_rejection,
    system_prompt,
)
from wordle_five.reward import Episode, metrics, score  # noqa: E402

TRACES = Path(__file__).parent.parent / "public" / "traces"
OUT = TRACES / "wordle"
ENDPOINT = os.environ.get("PIG_DEMO_INFERENCE", "http://100.127.247.67:8001/v1/chat/completions")
MODEL = os.environ.get("PIG_DEMO_MODEL", "brain-qwen38-dspark")

ARMS = {
    "base-off": {"label": "Out of the box", "thinking": False},
    "base-on": {"label": "Allowed to think", "thinking": True},
    "solver": {"label": "Best-known play", "thinking": None},
    "cautious": {"label": "Never wastes a guess", "thinking": None},
}


def call_model(messages: list[dict], thinking: bool, max_tokens: int | None = None, timeout: int = 300) -> dict:
    """One completion. Returns reply, reasoning and the real call metrics.

    Never raises on an upstream failure — a dropped call becomes a turn with a
    null reply, which the game scores as a rejected guess. A capture that
    silently discarded failed turns would be reporting a better model than the
    one that ran.
    """
    body = {
        "model": MODEL,
        "messages": messages,
        "temperature": 0.7,
        "max_tokens": max_tokens or (2048 if thinking else 512),
        "chat_template_kwargs": {"enable_thinking": bool(thinking)},
    }
    request = urllib.request.Request(
        ENDPOINT,
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"},
    )
    started = time.time()
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            payload = json.loads(response.read())
    except (urllib.error.URLError, TimeoutError, OSError) as exc:
        return {
            "reply": None,
            "reasoning": None,
            "call": {
                "promptTokens": None,
                "completionTokens": None,
                "reasoningTokens": None,
                "durationMs": round((time.time() - started) * 1000),
                "finishReason": f"error: {type(exc).__name__}",
            },
        }

    elapsed = round((time.time() - started) * 1000)
    choice = payload["choices"][0]
    message = choice.get("message", {})
    usage = payload.get("usage", {}) or {}
    details = usage.get("completion_tokens_details") or {}
    return {
        "reply": message.get("content"),
        "reasoning": message.get("reasoning_content"),
        "call": {
            "promptTokens": usage.get("prompt_tokens"),
            "completionTokens": usage.get("completion_tokens"),
            "reasoningTokens": details.get("reasoning_tokens"),
            "durationMs": elapsed,
            "finishReason": choice.get("finish_reason"),
        },
    }


def cautious_turn(game: Game) -> dict:
    """Only ever guesses a word that could still be the answer.

    This arm exists because the reward's counterweight is a claim about a
    trade-off, and a trade-off with only one policy on the board is a claim
    nobody can check. It never spends a turn on a word that cannot win, so it
    takes `consistency` outright — and it pays for that in turns, because it
    cannot buy information with a guess that has no chance. Weight the reward
    one way and it beats the entropy solver; weight it the other and it loses.
    That is the whole argument, made with two recorded runs instead of a claim.
    """
    pool = answers()
    alive = [w for w in S.consistent_candidates(game.history) if w not in {g for g, _ in game.history}]
    if not alive:
        guess = "tares"
    else:
        # Most informative CANDIDATE, not first alphabetically. Same objective as
        # the solver, restricted action set — which is the honest contrast. A
        # policy that opens on whatever sorts first looks incompetent rather than
        # cautious, and would make the trade-off it exists to demonstrate look
        # like a straw man.
        index = np.array([pool.index(w) for w in alive])
        best, best_bits = alive[0], -1.0
        for word in alive:
            counts: dict[int, int] = {}
            row = S.pattern_matrix()[pool.index(word)]
            for j in index:
                code = int(row[j])
                counts[code] = counts.get(code, 0) + 1
            total = len(alive)
            bits = -sum((c / total) * math.log2(c / total) for c in counts.values())
            if bits > best_bits:
                best, best_bits = word, bits
        guess = best
    return {
        "reply": f"[{guess}]",
        "reasoning": None,
        "call": {
            "promptTokens": None,
            "completionTokens": None,
            "reasoningTokens": None,
            "durationMs": None,
            "finishReason": "generated",
        },
    }


def solver_turn(game: Game) -> dict:
    """The reference player, recorded in the same shape as a model turn.

    durationMs is null rather than invented: nothing waited for this, and the
    player must not pretend otherwise. The UI renders a null duration as an
    instant step and labels the run as generated.
    """
    pool = answers()
    if not game.history:
        guess = pool[S._best_opener()]
    else:
        alive = S.consistent_candidates(game.history)
        guess = (
            pool[S.best_guess(np.array([pool.index(w) for w in alive]))] if alive else "tares"
        )
    return {
        "reply": f"[{guess}]",
        "reasoning": None,
        "call": {
            "promptTokens": None,
            "completionTokens": None,
            "reasoningTokens": None,
            "durationMs": None,
            "finishReason": "generated",
        },
    }


def capture(arm: str, seed: int) -> dict:
    config = ARMS[arm]
    game = Game(seed=seed)
    messages = [
        {"role": "system", "content": system_prompt()},
        {"role": "user", "content": "Enter your guess to begin."},
    ]
    turns: list[dict] = []

    while not game.over and len(turns) < MAX_GUESSES * 2:
        if arm == "solver":
            result = solver_turn(game)
        elif arm == "cautious":
            result = cautious_turn(game)
        else:
            result = call_model(messages, bool(config["thinking"]))

        guess = parse_guess(result["reply"])
        left = MAX_GUESSES - len(game.history)

        if guess is None:
            game.rejected += 1
            observation = render_rejection("no bracketed guess found", left)
        else:
            pattern, rejection = game.play(guess)
            observation = (
                render_rejection(rejection, left)
                if rejection is not None
                else render_feedback(guess, pattern or "", MAX_GUESSES - len(game.history))
            )

        turns.append({**result, "info": {"guess": guess, "observation": observation}})
        messages.append({"role": "assistant", "content": result["reply"] or ""})
        messages.append({"role": "user", "content": observation})

        if game.rejected >= MAX_GUESSES:
            break

    episode = Episode(
        answer=game.answer,
        guesses=[g for g, _ in game.history],
        patterns=[p for _, p in game.history],
        rejected=game.rejected,
        reference_depth=S.reference_depth(game.answer),
    )
    # A run that never reached a terminal state is marked truncated, and the
    # browser renders its verification as "unverifiable" rather than as a zero.
    truncated = not game.over and game.rejected < MAX_GUESSES

    return {
        "runId": f"{arm}-s{seed}",
        "seed": seed,
        "model": {"solver": "entropy-solver", "cautious": "candidate-only-solver"}.get(arm, MODEL),
        "capturedAt": time.strftime("%Y-%m-%d"),
        "rewards": score(episode),
        "metrics": metrics(episode),
        "truncated": truncated,
        "outcome": "solved" if game.solved else "failed",
        "answer": game.answer,
        "turns": turns,
    }


# ------------------------------------------------------------ alert-triage --
#
# The same fixture shape as wordle — runId, seed, model, capturedAt, rewards,
# metrics, truncated, outcome, turns[{reply, reasoning, call, info}] — driven
# through `alert_triage.taskset.play_episode`, the loop the probe and the
# tests share. The browser regenerates the world from the seed and replays the
# reply strings; nothing else in the fixture is trusted by the page.

AT_ARMS = {
    "base-off": {"label": "Out of the box", "thinking": False, "max_tokens": 1024, "timeout": 300},
    "base-on": {"label": "Allowed to think", "thinking": True, "max_tokens": 4096, "timeout": 900},
    "fast": {"label": "Reads the screen", "thinking": None, "policy": "fast"},
    "targeted": {"label": "Checks the hidden tells", "thinking": None, "policy": "targeted"},
    "thorough": {"label": "Runs the full procedure", "thinking": None, "policy": "thorough"},
}
AT_MODEL_NAME = {"fast": "fast-analyst", "targeted": "targeted-analyst", "thorough": "thorough-analyst"}

# A thinking budget the model exhausts is recorded as finishReason "length"
# with whatever content survived (usually none, which the engine rejects).
# That is a thing the model did under the budget it was given, not a capture
# error, and the fixture says so rather than retrying until it looks better.


def _messages_from(prompt: str, transcript: list[dict]) -> list[dict]:
    """The chat a model sees: the system prompt, the screen, then each reply and what it got back."""
    messages = [{"role": "system", "content": prompt}, {"role": "user", "content": transcript[0]["observation"]}]
    for entry in transcript[1:]:
        messages.append({"role": "assistant", "content": entry["reply"] or ""})
        messages.append({"role": "user", "content": entry["observation"]})
    return messages


def _generated_call() -> dict:
    return {"promptTokens": None, "completionTokens": None, "reasoningTokens": None, "durationMs": None, "finishReason": "generated"}


def capture_alert_triage(arm: str, seed: int) -> dict:
    config = AT_ARMS[arm]
    calls: list[dict] = []

    if config["thinking"] is None:
        policy = AT.POLICIES[config["policy"]]

        def respond(prompt: str, transcript: list[dict], view: dict) -> str | None:
            reply = policy(view)
            calls.append({"reply": reply, "reasoning": None, "call": _generated_call()})
            return reply
    else:
        def respond(prompt: str, transcript: list[dict], view: dict) -> str | None:
            result = call_model(_messages_from(prompt, transcript), bool(config["thinking"]),
                                max_tokens=config["max_tokens"], timeout=config["timeout"])
            calls.append(result)
            return result["reply"]

    world = world_for_seed(seed)
    played = AT_taskset.play_episode(seed, respond, world)
    steps = played["transcript"][1:]
    assert len(steps) == len(calls), "one model call per engine step"

    turns = [
        {
            **call,
            "info": {"action": step["action"], "rejection": step["rejection"], "observation": step["observation"]},
        }
        for call, step in zip(calls, steps)
    ]
    return {
        "runId": f"{arm}-s{seed}",
        "seed": seed,
        "model": AT_MODEL_NAME.get(arm, MODEL),
        "capturedAt": time.strftime("%Y-%m-%d"),
        "rewards": played["rewards"],
        "metrics": played["metrics"],
        "truncated": played["truncated"],
        "outcome": played["outcome"],
        "info": {**played["info"], "tier": world["tier"], "screen": played["transcript"][0]["observation"]},
        "turns": turns,
    }


def _at_summary(episode: dict) -> str:
    r = episode["rewards"]
    fmt = lambda v: "  -- " if v is None else f"{v:.2f}"  # noqa: E731
    actions = []
    for t in episode["turns"]:
        a = t["info"]["action"]
        if a is None:
            actions.append("REJ")
        elif a["action"] == "lookup":
            actions.append(a.get("month") or a.get("id") or a["what"])
        else:
            actions.append(a["action"].upper())
    return (
        f"{episode['info']['tier']:<8}{episode['info']['template']:<3} {episode['outcome']:<8}"
        f" caught={fmt(r['caught'])} hours={fmt(r['hours'])} evid={fmt(r['evidence'])}"
        f"  {episode['metrics']['hours_spent']:.2f}h  {actions}"
    )


TASKSETS = {
    "wordle-five": {"arms": ARMS, "out": TRACES / "wordle", "capture": capture},
    "alert-triage": {"arms": AT_ARMS, "out": TRACES / "alert-triage", "capture": capture_alert_triage},
}


def parse_seeds(spec: str) -> list[int]:
    """"0-7", "1,3", or a mix: "1-7,9"."""
    seeds: list[int] = []
    for part in spec.split(","):
        if "-" in part:
            lo, hi = part.split("-")
            seeds.extend(range(int(lo), int(hi) + 1))
        else:
            seeds.append(int(part))
    return seeds


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--taskset", default="wordle-five", choices=sorted(TASKSETS))
    parser.add_argument("--arm", required=True)
    parser.add_argument("--seeds", default="0-7")
    args = parser.parse_args()

    taskset = TASKSETS[args.taskset]
    if args.arm not in taskset["arms"]:
        parser.error(f"--arm must be one of {sorted(taskset['arms'])} for {args.taskset}")

    out: Path = taskset["out"]
    out.mkdir(parents=True, exist_ok=True)
    for seed in parse_seeds(args.seeds):
        if args.taskset == "alert-triage" and is_held_out(seed):
            # The held-out bucket is never captured, probed or trained on. A
            # fixture for one would put a held-out alert behind a permalink.
            print(f"{args.arm}-s{seed}: seed {seed} is held out — skipped", flush=True)
            continue
        started = time.time()
        episode = taskset["capture"](args.arm, seed)
        path = out / f"{episode['runId']}.json"
        path.write_text(json.dumps(episode, indent=2) + "\n")
        if args.taskset == "wordle-five":
            guesses = [t["info"]["guess"] for t in episode["turns"]]
            print(
                f"{episode['runId']:>16}  {episode['answer']}  {episode['outcome']:<7}"
                f"  solved={episode['rewards']['solved']:.0f}"
                f" econ={episode['rewards']['economy']:.2f}"
                f" cons={episode['rewards']['consistency']:.2f}"
                f"  {time.time()-started:5.1f}s  {guesses}",
                flush=True,
            )
        else:
            print(f"{episode['runId']:>16}  {_at_summary(episode)}  {time.time()-started:6.1f}s", flush=True)
    return 0


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