"""Reliability harness: measure how cleanly a served model speaks the action schema.

Runs the REAL agent prompts (ttw.agents.build_messages over a seeded world for a
few turns) against a backend, parses each raw response through the same layer the
game uses, and reports the metrics that quantify the Well-Tuned win:

    valid_json_pct      fraction of responses that yield a JSON object
    valid_offer_pct     fraction of raw offers that coerce to a real Offer
    self_buy_rate       fraction of BUY offers that buy the creature's own product
    avg_offers_per_turn mean valid offers per creature-turn

Workflow for the badge: deploy the base model, eval; deploy the fine-tuned model
(TTW_MODEL=...), eval; compare the two tables in Field Notes.

Backends:
    --backend modal   call the currently deployed ttw-serve endpoint (default)
    --backend local   run a local transformers pipeline on --model (needs a GPU/CPU + deps)
    --backend dummy    no network: drive the deterministic dummy policy (verifies harness logic)

Usage:
    python scripts/finetune/eval_model.py                       # eval the live endpoint
    python scripts/finetune/eval_model.py --backend dummy       # self-test, no GPU
    python scripts/finetune/eval_model.py --backend local --model AdmiralTaco/ttw-trader-0.5b
"""

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from ttw.actions import _coerce_offer, _extract_json  # noqa: E402
from ttw.agents import build_messages  # noqa: E402
from ttw.world import seed_world  # noqa: E402


def collect_modal(turns: int, max_tokens: int, temperature: float) -> list[tuple[str, str]]:
    """Drive the deployed ttw-serve endpoint. Returns [(creature, raw_text), ...]."""
    from ttw.llm import ModalLLM
    from ttw.sim import step

    client = ModalLLM()
    samples: list[tuple[str, str]] = []
    world = seed_world()

    # Use a real LLM policy so the world evolves between turns the same way the
    # game does, while we snapshot each turn's raw text for scoring.
    from ttw.agents import make_llm_policy

    policy = make_llm_policy(client, temperature=temperature, max_tokens=max_tokens)
    for _ in range(turns):
        step(world, policy)
        for name, raw in policy.state["raw"].items():
            samples.append((name, raw))
    return samples


def collect_local(model_id: str, turns: int, max_tokens: int, temperature: float) -> list[tuple[str, str]]:
    """Run a local transformers chat pipeline over the real prompts."""
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tok = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)

    samples: list[tuple[str, str]] = []
    world = seed_world()

    def local_chat(messages: list[dict]) -> str:
        prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = tok(prompt, return_tensors="pt").to(model.device)
        out = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=temperature,
            do_sample=temperature > 0,
            pad_token_id=tok.eos_token_id,
        )
        return tok.decode(out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True)

    from ttw.actions import parse_actions
    from ttw.sim import step

    def policy(w, name):
        msgs = build_messages(w, w.creatures[name])
        raw = local_chat(msgs)
        samples.append((name, raw))
        return parse_actions(name, raw)

    for _ in range(turns):
        step(world, policy)
    return samples


def collect_dummy(turns: int) -> list[tuple[str, str]]:
    """No-network path: synthesize raw JSON from the dummy policy's offers.

    Lets the scoring logic be exercised end-to-end without any model. The dummy
    never self-buys, so self_buy_rate should read 0.0, a useful sanity check.
    """
    from ttw.dummy import make_random_policy
    from ttw.sim import step

    base_policy = make_random_policy(seed=7)
    samples: list[tuple[str, str]] = []

    def policy(w, name):
        offers, gossip = base_policy(w, name)
        raw = json.dumps(
            {
                "thought": f"{name} works the market",
                "offers": [
                    {"side": o.side, "good": o.good, "price": o.price, "qty": o.qty}
                    for o in offers
                ],
                "gossip": gossip[0].message if gossip else "",
            }
        )
        samples.append((name, raw))
        return offers, gossip

    world = seed_world()
    for _ in range(turns):
        step(world, policy)
    return samples


def score(samples: list[tuple[str, str]]) -> dict:
    """Compute reliability metrics from (creature, raw_text) samples."""
    produces = {name: c.produces for name, c in seed_world().creatures.items()}

    n_turns = len(samples)
    valid_json = 0
    raw_offer_count = 0
    valid_offer_count = 0
    buy_count = 0
    self_buy_count = 0

    for creature, raw in samples:
        payload = _extract_json(raw)
        if payload is None:
            continue
        valid_json += 1
        raw_offers = payload.get("offers", [])
        if isinstance(raw_offers, dict):
            raw_offers = [raw_offers]
        if not isinstance(raw_offers, list):
            raw_offers = []
        for raw_offer in raw_offers:
            if not isinstance(raw_offer, dict):
                continue
            raw_offer_count += 1
            offer = _coerce_offer(creature, raw_offer)
            if offer is None:
                continue
            valid_offer_count += 1
            if offer.side == "buy":
                buy_count += 1
                if offer.good == produces.get(creature):
                    self_buy_count += 1

    return {
        "samples": n_turns,
        "valid_json_pct": valid_json / n_turns if n_turns else 0.0,
        "valid_offer_pct": valid_offer_count / raw_offer_count if raw_offer_count else 0.0,
        "self_buy_rate": self_buy_count / buy_count if buy_count else 0.0,
        "avg_offers_per_turn": valid_offer_count / n_turns if n_turns else 0.0,
    }


def print_table(label: str, m: dict) -> None:
    """Print one clean metrics block."""
    print(f"\n=== {label} ({m['samples']} creature-turns) ===")
    print(f"  valid_json_pct      {m['valid_json_pct']:.1%}")
    print(f"  valid_offer_pct     {m['valid_offer_pct']:.1%}")
    print(f"  self_buy_rate       {m['self_buy_rate']:.1%}")
    print(f"  avg_offers_per_turn {m['avg_offers_per_turn']:.2f}")


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--backend", choices=["modal", "local", "dummy"], default="modal")
    ap.add_argument("--model", default=None, help="model id for --backend local")
    ap.add_argument("--turns", type=int, default=6)
    ap.add_argument("--max-tokens", type=int, default=256)
    ap.add_argument("--temperature", type=float, default=0.7)
    args = ap.parse_args()

    if args.backend == "modal":
        samples = collect_modal(args.turns, args.max_tokens, args.temperature)
        label = "modal:ttw-serve"
    elif args.backend == "local":
        if not args.model:
            ap.error("--backend local requires --model")
        samples = collect_local(args.model, args.turns, args.max_tokens, args.temperature)
        label = f"local:{args.model}"
    else:
        samples = collect_dummy(args.turns)
        label = "dummy"

    print_table(label, score(samples))


if __name__ == "__main__":
    main()