#!/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
"""

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"))

import numpy as np  # 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

OUT = Path(__file__).parent.parent / "public" / "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) -> 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": 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=300) 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,
    }


def parse_seeds(spec: str) -> list[int]:
    if "-" in spec:
        lo, hi = spec.split("-")
        return list(range(int(lo), int(hi) + 1))
    return [int(s) for s in spec.split(",")]


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

    OUT.mkdir(parents=True, exist_ok=True)
    for seed in parse_seeds(args.seeds):
        started = time.time()
        episode = capture(args.arm, seed)
        path = OUT / f"{episode['runId']}.json"
        path.write_text(json.dumps(episode, indent=2) + "\n")
        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}"
        )
    return 0


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