#!/usr/bin/env python3
"""
Evaluate one or more Tinker SFT checkpoints WITHOUT renting a GPU.

Samples each checkpoint on Tinker's GPUs, renders the generated Three.js
locally (headless Chromium / swiftshader, CPU), and duels each render against
the ARCHIVED REAL WINNER for that held-out prompt via a vision judge on
OpenRouter. Reports win-rate-vs-archived-winner per checkpoint plus a
first-vs-last delta, so you can see whether training actually improved renders
(not just lowered loss).

  - --save-every checkpoints are `save_state` (resume) checkpoints; we
    load_state() then save_weights_and_get_sampling_client() to sample them.
  - Multiple --checkpoint LABEL=tinker://... are compared on the SAME prompts.
    Archived-winner renders are rendered once and cached across checkpoints.

Judge model: the live validator uses GLM-4.1V-9B-Thinking, which OpenRouter has
retired. Default here is z-ai/glm-4.5v (nearest available) — a PROXY good for
relative comparison, not the exact validator verdict. Override with --judge-model.

Usage:
    python tinker/eval_tinker.py \
        --checkpoint first=tinker://.../weights/threejs-sft-v1-step25 \
        --checkpoint last=tinker://.../weights/threejs-sft-v1-step319 \
        --n-prompts 8
"""
from __future__ import annotations

import argparse
import asyncio
import base64
import json
import os
import random
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))                       # tinker_data (avoid SDK name clash)
sys.path.insert(0, str(HERE.parent / "eval"))       # self_judge_eval + judge_challenge

import tinker  # noqa: E402
from tinker import types as tt  # noqa: E402

from tinker_data import image_expected_tokens, _texts  # noqa: E402
from judge_challenge import evaluate_duel  # noqa: E402
from self_judge_eval import render_js, png_to_data_url, chat_text  # noqa: E402

try:
    import wandb
except ImportError:
    wandb = None


def read_line(path: str, line: int) -> str | None:
    try:
        with open(os.path.expanduser(path)) as f:
            lines = [ln.strip() for ln in f if ln.strip()]
        return lines[line] if line < len(lines) else None
    except OSError:
        return None


def build_prompt_input(row, tok, root):
    """Prompt-only ModelInput matching the training prefix EXACTLY (minus the
    completion): [system + nl<image>nl] [ImageChunk] [nl+user+nl]."""
    sys_txt, user_txt = _texts(row)
    img_path = os.path.join(root, row["images"][0])
    itoks = image_expected_tokens(img_path)
    img_bytes = open(img_path, "rb").read()
    pre = tok.encode(sys_txt + "\n<image>\n", add_special_tokens=False)
    post = tok.encode("\n" + user_txt + "\n", add_special_tokens=False)
    mi = tinker.ModelInput.empty()
    for t in pre:
        mi = mi.append_int(t)
    mi = mi.append(tt.ImageChunk(data=img_bytes, format="png", expected_tokens=itoks, type="image"))
    for t in post:
        mi = mi.append_int(t)
    return mi


def load_eval_rows(preference_jsonl, n_prompts):
    rows = [json.loads(l) for l in open(preference_jsonl) if l.strip()]
    rows = [r for r in rows if r.get("split") == "eval"]
    return rows[:n_prompts] if n_prompts else rows


def make_sampling_client(svc, label, state_path, base_model, rank):
    """load_state(step-N resume checkpoint) -> export a sampler on Tinker."""
    tc = svc.create_lora_training_client(base_model=base_model, rank=rank)
    print(f"[{label}] load_state: {state_path}", flush=True)
    tc.load_state(state_path).result()
    print(f"[{label}] exporting sampler (on Tinker)...", flush=True)
    sc = tc.save_weights_and_get_sampling_client(name=f"eval-{label}-sampler")
    tok = tc.get_tokenizer()
    return sc, tok


def generate_js(sc, tok, row, root, max_tokens, temperature, seed):
    mi = build_prompt_input(row, tok, root)
    resp = sc.sample(
        prompt=mi, num_samples=1,
        sampling_params=tinker.SamplingParams(max_tokens=max_tokens, temperature=temperature, seed=seed),
    ).result()
    toks = resp.sequences[0].tokens
    return tok.decode(toks, skip_special_tokens=True) if toks else None


async def main_async(args):
    # ---- renderer setup (single sidecar; randomize ports to dodge stragglers) ----
    sys.path.insert(0, args.pipeline_service_dir)
    from modules.renderer.module import RendererModule
    from config.settings import settings as ps
    ps.renderer.sidecar_count = 1
    ps.renderer.pool_size = 1
    ps.renderer.startup_timeout_s = 180.0
    base = random.randint(20, 200)
    ps.renderer.sidecar_port = 8003 + base
    ps.renderer.static_port_base = 3000 + base

    rows = load_eval_rows(args.preference_jsonl, args.n_prompts)
    print(f"Loaded {len(rows)} held-out eval rows | checkpoints: {[c[0] for c in args.checkpoints]}", flush=True)
    if not rows:
        print("No eval rows — check --preference-jsonl / split=='eval'."); return

    # ---- judge client (OpenRouter) ----
    key = os.environ.get(args.judge_api_key_env) or read_line(args.judge_api_key_file, 0)
    if not key:
        raise SystemExit(f"No judge key: set ${args.judge_api_key_env} or put it in {args.judge_api_key_file}")
    from openai import AsyncOpenAI
    judge = AsyncOpenAI(base_url=args.judge_base_url, api_key=key)

    svc = tinker.ServiceClient()
    renderer = RendererModule(ps.renderer)
    await renderer.startup()

    if wandb is not None and args.wandb:
        wandb.init(project=args.wandb_project, group="eval", name=args.run_name, job_type="eval")

    results = {}   # label -> dict(wins,losses,draws,errors)
    archived_png_cache = {}  # prompt image_path -> png bytes (render once)

    try:
        for label, state_path in args.checkpoints:
            print(f"\n===== checkpoint '{label}' =====", flush=True)
            sc, tok = make_sampling_client(svc, label, state_path, args.base_model, args.rank)
            wins = losses = draws = errors = empty = 0
            for i, row in enumerate(rows):
                stem = row.get("name", f"eval_{i}")
                image_path = os.path.join(args.root, row["images"][0])
                try:
                    cand_js = generate_js(sc, tok, row, args.root, args.max_tokens, args.temperature, args.seed)
                    if not cand_js:
                        print(f"  [{label}][{stem}] empty gen"); errors += 1; continue
                    cand_png = await render_js(renderer, f"{label}_{stem}_cand", image_path, cand_js)
                    if image_path not in archived_png_cache:
                        arch_js = chat_text(row["chosen"][0])
                        archived_png_cache[image_path] = await render_js(renderer, f"{stem}_arch", image_path, arch_js)
                    arch_png = archived_png_cache[image_path]
                    if not cand_png or not arch_png:
                        print(f"  [{label}][{stem}] empty render (cand={bool(cand_png)} arch={bool(arch_png)})")
                        empty += 1; errors += 1; continue
                    with open(image_path, "rb") as f:
                        prompt_url = "data:image/png;base64," + base64.b64encode(f.read()).decode()
                    res = await evaluate_duel(
                        client=judge, prompt_url=prompt_url,
                        left_url=png_to_data_url(arch_png),   # archived winner
                        right_url=png_to_data_url(cand_png),  # our checkpoint
                        seed=42, model=args.judge_model,
                    )
                    if res.outcome > 0: wins += 1
                    elif res.outcome < 0: losses += 1
                    else: draws += 1
                    if args.save_renders:
                        d = Path(args.save_renders); d.mkdir(parents=True, exist_ok=True)
                        (d / f"{label}_{stem}_cand.png").write_bytes(cand_png)
                    n = wins + losses + draws
                    print(f"  [{label}][{n}/{len(rows)}] {stem}: outcome={res.outcome} "
                          f"(W{wins}/L{losses}/D{draws}) — {res.issues[:50]}", flush=True)
                except Exception as e:
                    print(f"  [{label}][{stem}] ERROR: {e}"); errors += 1
            nj = wins + losses + draws
            wr = wins / nj if nj else 0.0
            results[label] = dict(wins=wins, losses=losses, draws=draws, errors=errors, empty=empty, win_rate=wr, n=nj)
            print(f"[{label}] {wins}W/{losses}L/{draws}D/{errors}err ({empty} empty) — win_rate={wr:.1%}", flush=True)
    finally:
        await renderer.shutdown()

    # ---- comparison ----
    print("\n================ SUMMARY ================", flush=True)
    for label, r in results.items():
        print(f"  {label:>8}: win_rate={r['win_rate']:.1%}  ({r['wins']}W/{r['losses']}L/{r['draws']}D, "
              f"{r['errors']} err, n={r['n']})", flush=True)
    labels = [c[0] for c in args.checkpoints]
    if len(labels) >= 2:
        a, b = labels[0], labels[-1]
        d = results[b]["win_rate"] - results[a]["win_rate"]
        print(f"\n  Δ ({b} - {a}) = {d:+.1%} win-rate vs archived winners "
              f"→ training {'HELPED' if d > 0 else 'did NOT help' if d < 0 else 'unchanged'}", flush=True)
    total_empty = sum(r["empty"] for r in results.values())
    total_n = sum(r["n"] for r in results.values())
    if total_n == 0 and total_empty:
        print("  ALL renders empty — RENDERER problem, numbers meaningless.", flush=True)
    print("  (judge = proxy; win_rate>50% = beats archived REAL winners on held-out prompts)", flush=True)

    if wandb is not None and args.wandb:
        wandb.log({f"eval/{lab}/win_rate": r["win_rate"] for lab, r in results.items()})
        wandb.finish()


def parse_ckpt(s):
    if "=" not in s:
        raise argparse.ArgumentTypeError("use LABEL=tinker://...  e.g. last=tinker://...step319")
    lab, path = s.split("=", 1)
    return (lab, path)


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--checkpoint", dest="checkpoints", action="append", type=parse_ckpt, required=True,
                   help="LABEL=tinker://...  (repeatable; compared on the same prompts)")
    p.add_argument("--base-model", default="Qwen/Qwen3.6-27B")
    p.add_argument("--rank", type=int, default=32)
    p.add_argument("--preference-jsonl", default="datasets_dpo/preference.jsonl")
    p.add_argument("--root", default=".")
    p.add_argument("--n-prompts", type=int, default=8)
    p.add_argument("--max-tokens", type=int, default=16384)
    p.add_argument("--temperature", type=float, default=0.6)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--save-renders", default=None, help="dir to dump candidate PNGs for inspection")

    p.add_argument("--pipeline-service-dir",
                   default=str(Path(__file__).resolve().parents[2] / "miner-submission" / "pipeline_service"))

    p.add_argument("--judge-base-url", default="https://openrouter.ai/api/v1")
    p.add_argument("--judge-model", default="z-ai/glm-4.5v",
                   help="Validator uses GLM-4.1V-9B-Thinking (retired on OpenRouter); this is a proxy.")
    p.add_argument("--judge-api-key-env", default="OPENROUTER_API_KEY")
    p.add_argument("--judge-api-key-file", default=os.path.expanduser("~/openrouterkey"))
    p.add_argument("--api-key-file", default="/Users/xavier/bittensor/affine/package_finetune32B/api_key.md",
                   help="line 1 = Tinker key (for TINKER_API_KEY if unset).")

    p.add_argument("--wandb", action="store_true")
    p.add_argument("--wandb-project", default="404-coder-tinker")
    p.add_argument("--run-name", default="eval-compare")
    args = p.parse_args()

    if "TINKER_API_KEY" not in os.environ:
        k = read_line(args.api_key_file, 0)
        if k:
            os.environ["TINKER_API_KEY"] = k
    asyncio.run(main_async(args))


if __name__ == "__main__":
    main()