#!/usr/bin/env python3
"""Replay the retained Muse/Flash VLM comparison with GPT-6 Astra xhigh.

Uses muse_judge_pilot._parts_for and the archived prompt/schema/headers.
Defaults to Astra; --provider openrouter reuses this runner for other judges.
Existing judges, images, records and DB rows are preserved.
Each order is an independent request; --live is required to send requests.
"""
from __future__ import annotations

import argparse
import base64
import hashlib
import json
import os
import statistics
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path

import httpx
from dotenv import load_dotenv
from jsonschema import validate
from openai import OpenAI

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

MODEL = "gpt-6-astra"
EFFORT = "xhigh"
LABELS = ["A", "B"]
BASELINE = ROOT / "artifact/20260903_muse_judge_pilot"
DEFAULT_OUT = ROOT / "artifact/20260905_astra_xhigh_vlm_compare"


def digest(value):
    data = value if isinstance(value, bytes) else json.dumps(
        value, ensure_ascii=False, sort_keys=True, separators=(",", ":")
    ).encode()
    return hashlib.sha256(data).hexdigest()


def save(path, data):
    path.parent.mkdir(parents=True, exist_ok=True)
    temp = path.with_suffix(path.suffix + ".tmp")
    temp.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    temp.replace(path)


def manifest_parts(parts):
    output = []
    for part in parts:
        if part["type"] == "text":
            output.append(part)
        elif part["type"] == "image_url":
            image = part["image_url"]
            header, encoded = image["url"].split(",", 1)
            data = base64.b64decode(encoded, validate=True)
            output.append({"type": "image_url", "mime": header[5:].split(";")[0],
                           "sha256": digest(data), "bytes": len(data),
                           "detail": image.get("detail", "auto")})
        else:
            raise ValueError(f"Unsupported retained part: {part['type']}")
    return output


def ask(job, key, system, schema, max_tokens, model=MODEL, effort=EFFORT,
        provider="openai", single=False):
    if provider == "gemini":
        from gemini_pro_judge_pilot import ask as native_ask
        return native_ask(job, system, schema, single=single, max_tokens=max_tokens)
    from app.modules.pipeline.multiroll_select import normalize_flip_verdict

    tag, order, parts, request_identity = job
    mapping = {} if single else dict(zip(LABELS, LABELS if order == "forward" else reversed(LABELS)))
    base_url = "https://openrouter.ai/api/v1" if provider == "openrouter" else "https://api.openai.com/v1"
    slot = {"model": model, "reasoning_effort": effort, "provider": provider, "order": order,
            "display_to_canonical": mapping, "request_identity": request_identity,
            "started_at": datetime.now(timezone.utc).isoformat(), "ok": False}
    attempts = []

    def sent(request):
        attempts.append({"method": request.method, "path": request.url.path})

    started = time.monotonic()
    try:
        options = ({"max_tokens": max_tokens, "extra_body": {
            "reasoning": {"effort": effort}, "provider": {"require_parameters": True}}}
            if provider == "openrouter" else {
                "reasoning_effort": effort, "max_completion_tokens": max_tokens, "store": False})
        slot["request_options"] = options
        with OpenAI(api_key=key, base_url=base_url,
                    max_retries=0, timeout=300,
                    http_client=httpx.Client(event_hooks={"request": [sent]})) as client:
            raw = client.chat.completions.with_raw_response.create(
                model=model,
                messages=[{"role": "system", "content": system},
                          {"role": "user", "content": parts}],
                response_format={"type": "json_schema", "json_schema": {
                    "name": "retained_still_judge", "strict": True, "schema": schema}},
                **options,
            )
            response = raw.parse()
            slot.update(http_status=raw.status_code, request_id=raw.headers.get("x-request-id"),
                        response_id=response.id, returned_model=response.model,
                        usage=response.usage.model_dump(mode="json") if response.usage else None,
                        finish_reason=response.choices[0].finish_reason,
                        response=response.model_dump(mode="json"))
            if response.choices[0].finish_reason != "stop":
                raise ValueError(f"Incomplete response: {slot['finish_reason']}")
            verdict = json.loads(response.choices[0].message.content or "")
            validate(verdict, schema)
            slot["raw"] = verdict
            slot["normalized"] = verdict if single else normalize_flip_verdict(verdict, mapping, LABELS)
            slot["ok"] = True
    except Exception as exc:
        slot.update(error_type=type(exc).__name__, error=str(exc).replace(key, "[REDACTED]"))
        if getattr(exc, "status_code", None):
            slot["http_status"] = exc.status_code
    slot["elapsed_seconds"] = round(time.monotonic() - started, 3)
    slot["http_attempts"] = attempts
    slot["http_attempt_count"] = len(attempts)
    return tag, order, slot


def compare(results, records, tags):
    from app.modules.pipeline.multiroll_select import combine_flip_verdicts

    summary = {"note": "Selection agreement is not accuracy; no human ground-truth labels.",
               "historical_models_recalled": False, "shots": [], "models": {}}
    for model, per_tag in results.items():
        counts = {"ok_calls": 0, "errors": 0, "stable_pairs": 0, "complete_pairs": 0,
                  "matches_production": 0, "hard_violation_mentions": 0,
                  "all_candidates_fail_calls": 0}
        elapsed, tokens = [], {"input": 0, "output": 0, "reasoning": 0}
        for tag in tags:
            orders = per_tag.get(tag, {})
            for slot in orders.values():
                # A failed/truncated response still consumed time and tokens.
                if "elapsed_seconds" in slot:
                    elapsed.append(slot["elapsed_seconds"])
                usage = slot.get("usage") or {}
                tokens["input"] += usage.get("prompt_tokens", 0)
                tokens["output"] += usage.get("completion_tokens", 0)
                tokens["reasoning"] += (usage.get("completion_tokens_details") or {}).get("reasoning_tokens", 0)
                if not slot.get("ok"):
                    counts["errors"] += 1
                    continue
                counts["ok_calls"] += 1
                norm = slot["normalized"]
                counts["hard_violation_mentions"] += sum(
                    len(row.get("hard_violations", [])) for row in norm.get("readings", []))
                counts["all_candidates_fail_calls"] += bool(norm.get("all_candidates_fail"))
            if all(orders.get(o, {}).get("ok") for o in ("forward", "reverse")):
                fwd, rev = [orders[o]["normalized"] for o in ("forward", "reverse")]
                chosen, combined = combine_flip_verdicts(fwd, rev, LABELS, LABELS)
                counts["complete_pairs"] += 1
                counts["stable_pairs"] += fwd["winner"] == rev["winner"]
                counts["matches_production"] += chosen == records[tag].get("selected")
                summary["shots"].append({"tag": tag, "model": model,
                    "forward": fwd["winner"], "reverse": rev["winner"],
                    "combined": chosen, "score_totals": combined["totals"],
                    "production_selected": records[tag].get("selected")})
        if elapsed:
            counts.update(mean_seconds=round(statistics.mean(elapsed), 3),
                          median_seconds=round(statistics.median(elapsed), 3), tokens=tokens)
        summary["models"][model] = counts
    return summary


def single_contract(original, pair_schema):
    """Retain the rubric, replacing only comparative framing/output instructions."""
    import copy

    start = original.index("LOOK BEFORE YOU SCORE.")
    end = original.index("\nA candidate that reads wrong on DIRECTION")
    rubric = original[start:end]
    old = ("A candidate that realises the shot text's\n"
           "   framing beats any candidate that shows MORE of the scene, MORE\n"
           "   of a body, or MORE contract items by widening or reframing.")
    if rubric.count(old) != 1:
        raise ValueError("Archived comparative framing changed")
    rubric = rubric.replace(old, "Follow the shot text's framing. Showing MORE of the scene, MORE\n"
                            "   of a body, or MORE contract items by widening or reframing\n"
                            "   does not improve fidelity.")
    system = ("You are given ONE image-generation prompt, its labelled REFERENCE images,\n"
              "and exactly ONE CANDIDATE image. Evaluate this candidate independently.\n"
              "There is no second candidate to compare, rank, or infer.\n\n" + rubric +
              "\nJudge the single candidate on the same 0-10 integer fidelity scale.\n"
              "A wrong DIRECTION, BUILT SPACE, or PHYSICS cannot be compensated by\n"
              "generic beauty; those axes do not trade against each other.\n"
              "Set candidate_fails if this candidate fails an axis and the shot needs\n"
              "to be taken again. Do not assume a better or worse unseen alternative.\n"
              "Output readings on all FOUR axes and hard_violations, then score,\n"
              "a one-line verdict_ko citing decisive prompt points, and candidate_fails.\n"
              "Write every prose field in Korean.")
    reading = copy.deepcopy(pair_schema["properties"]["readings"]["items"])
    reading["properties"].pop("label")
    reading["required"].remove("label")
    schema = {"type": "object", "additionalProperties": False,
              "properties": {"readings": reading,
                             "score": {"type": "integer", "minimum": 0, "maximum": 10},
                             "verdict_ko": {"type": "string"},
                             "candidate_fails": {"type": "boolean"}},
              "required": ["readings", "score", "verdict_ko", "candidate_fails"]}
    return system, schema


def run_single(args):
    """One candidate per request, independent identical repetitions, no DB writes."""
    import difflib
    from muse_judge_pilot import _parts_for

    baseline = json.loads((args.baseline / "input_manifest.json").read_text())
    records_path = Path(baseline["records"])
    records = json.loads(records_path.read_text())
    original = (args.baseline / "judge_sys.txt").read_text()
    pair_schema = json.loads((args.baseline / "judge_schema.json").read_text())
    headers = json.loads((args.baseline / "headers.json").read_text())
    system, schema = single_contract(original, pair_schema)
    tags = baseline["shots"][:args.limit] if args.limit else baseline["shots"]
    if args.repeats < 2:
        raise ValueError("Independent repeatability needs at least two repetitions")
    args.out.mkdir(parents=True, exist_ok=True)
    path = args.out / "results.json"
    results = json.loads(path.read_text()) if path.exists() else {args.model: {}}
    mine = results.setdefault(args.model, {})
    jobs, inputs = [], {}
    sources = {str(p): digest(p.read_bytes()) for p in (
        records_path, args.baseline / "results.json", args.baseline / "input_manifest.json",
        args.baseline / "judge_sys.txt", args.baseline / "judge_schema.json")}
    for tag in tags:
        paired = _parts_for(records[tag], records_path.parent, tag, LABELS, LABELS,
                            headers["default"], headers["bgfirst"])
        if manifest_parts(paired) != baseline["inputs"][f"{tag}_forward"]["parts"]:
            raise ValueError(f"Source input changed: {tag}")
        for candidate, index in (("A", -3), ("B", -1)):
            # Prefix is the same shot prompt and shared reference images; only
            # the comparative header and the two-candidate tail are replaced.
            prefix = paired[:-4]
            header = headers["bgfirst"] if records[tag].get("bgfirst") else headers["default"]
            if not prefix[0]["text"].startswith(header + "\n"):
                raise ValueError("Unexpected archived header")
            parts = [{"type": "text", "text": "THE PROMPT FOR THIS SINGLE CANDIDATE:\n" + records[tag]["prompt"]}]
            parts += prefix[1:] + [{"type": "text", "text": "CANDIDATE IMAGE — evaluate this image only:"}, paired[index]]
            retained = manifest_parts(parts)
            identity = digest({"model": args.model, "effort": args.effort,
                               "provider": args.provider, "max_tokens": args.max_tokens,
                               "system": system, "schema": schema, "parts": retained})
            inputs[f"{tag}_{candidate}"] = {"request_identity": identity, "parts": retained,
                                           "candidate": candidate, "reference_count": len(records[tag].get("refs", []))}
            for repeat in range(1, args.repeats + 1):
                key = f"{candidate}_r{repeat}"
                old = mine.get(tag, {}).get(key)
                if old:
                    if old.get("request_identity") != identity:
                        raise ValueError("Input drift: use a fresh single-image output directory")
                    # Errors are also durable observations, never silently retried.
                    continue
                jobs.append((tag, key, parts, identity))
    if len(jobs) > args.max_calls:
        raise ValueError(f"{len(jobs)} requests exceed max-calls={args.max_calls}")
    manifest = {"evaluation": "single_candidate", "model": args.model, "effort": args.effort,
                "provider": args.provider, "max_completion_tokens": args.max_tokens,
                "sdk_max_retries": 0, "repeats": args.repeats, "workers": args.workers,
                "shots": tags, "records": str(records_path), "project_id": args.project_id,
                "episode_id": args.episode_id, "baseline": str(args.baseline),
                "pending_calls": len(jobs), "inputs": inputs, "historical_source_sha256": sources,
                "system_sha256": digest(system.encode()), "schema_sha256": digest(schema),
                "reference_policy": "Original shared references retained; exactly one candidate, no A/B label, no other judgment.",
                "key_slot": args.key_slot}
    save(args.out / "input_manifest.json", manifest)
    (args.out / "judge_sys.txt").write_text(system)
    save(args.out / "judge_schema.json", schema)
    (args.out / "prompt.diff").write_text("".join(difflib.unified_diff(
        original.splitlines(keepends=True), system.splitlines(keepends=True),
        fromfile="paired_judge_sys", tofile="single_judge_sys")))
    print(json.dumps({"model": args.model, "effort": args.effort, "evaluation": "single",
                      "images": len(inputs), "repeats": args.repeats, "calls": len(jobs),
                      "workers": args.workers, "live": args.live}), flush=True)
    if not args.live:
        return 0
    key_name = ("OPENROUTER_API_KEY" if args.provider == "openrouter" else
                "OPENAI_API_KEY_SECONDARY" if args.key_slot == "secondary" else "OPENAI_API_KEY")
    key = "" if args.provider == "gemini" else os.environ.get(key_name, "").strip()
    if args.provider != "gemini" and not key:
        raise ValueError(f"{key_name} missing")
    started = time.monotonic()
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = [pool.submit(ask, job, key, system, schema, args.max_tokens,
                               args.model, args.effort, args.provider, True) for job in jobs]
        for future in as_completed(futures):
            tag, name, slot = future.result()
            candidate, repeat = name.split("_r")
            slot.update(candidate=candidate, repeat=int(repeat), evaluation="single_candidate")
            mine.setdefault(tag, {})[name] = slot
            save(path, results)
            print(json.dumps({"tag": tag, "candidate": candidate, "repeat": repeat,
                              "ok": slot["ok"], "score": slot.get("raw", {}).get("score"),
                              "seconds": slot["elapsed_seconds"], "error": slot.get("error")}, ensure_ascii=False), flush=True)
    slots = [s for rows in mine.values() for s in rows.values()]
    save(args.out / "run_summary.json", {"model": args.model, "calls": len(slots),
        "ok": sum(bool(s["ok"]) for s in slots), "http_attempts": sum(s["http_attempt_count"] for s in slots),
        "wall_seconds": round(time.monotonic()-started, 3),
        "historical_sources_unchanged": all(digest(Path(p).read_bytes()) == h for p, h in sources.items())})
    return 0 if all(s["ok"] for s in slots) else 1


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("project_id")
    parser.add_argument("episode_id")
    parser.add_argument("--baseline", type=Path, default=BASELINE)
    parser.add_argument("--out", type=Path, default=DEFAULT_OUT)
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--max-calls", type=int, default=12)
    parser.add_argument("--max-tokens", type=int, default=16000)
    parser.add_argument("--key-slot", choices=("primary", "secondary"), default="secondary")
    parser.add_argument("--model", default=MODEL)
    parser.add_argument("--effort", default=EFFORT)
    parser.add_argument("--provider", choices=("openai", "openrouter", "gemini"), default="openai")
    parser.add_argument("--single", action="store_true")
    parser.add_argument("--repeats", type=int, default=2)
    parser.add_argument("--live", action="store_true")
    args = parser.parse_args()
    if not 1 <= args.workers <= 8:
        parser.error("workers must be between 1 and 8")
    if args.out.resolve() == args.baseline.resolve():
        parser.error("output must not replace the historical comparison")
    load_dotenv(BACKEND / ".env", override=False)
    if args.single:
        return run_single(args)
    if args.provider == "gemini":
        parser.error("Use gemini_pro_judge_pilot for paired native Gemini replay")
    from muse_judge_pilot import _parts_for, _load_results
    from app.core.config import settings

    rdir = Path(settings.projects_dir) / args.project_id / "images" / args.episode_id / "scene/recipe"
    records_path = rdir / "records.json"
    records = json.loads(records_path.read_text(encoding="utf-8"))
    historical = _load_results(args.baseline / "results.json")
    model_tags = [set(rows) for rows in historical.values()]
    if not model_tags or any(tags != model_tags[0] for tags in model_tags):
        raise ValueError("Historical model populations differ")
    tags = sorted(model_tags[0])
    if args.limit:
        tags = tags[:args.limit]
    snapshots = {name: (args.baseline / name).read_bytes()
                 for name in ("judge_sys.txt", "judge_schema.json", "headers.json")}
    system = snapshots["judge_sys.txt"].decode()
    schema = json.loads(snapshots["judge_schema.json"])
    headers = json.loads(snapshots["headers.json"])
    original_hashes = {str(p): digest(p.read_bytes()) for p in
                       [records_path, args.baseline / "results.json"]}
    args.out.mkdir(parents=True, exist_ok=True)
    for name, data in snapshots.items():
        (args.out / name).write_bytes(data)
    results_path = args.out / "results.json"
    results = _load_results(results_path) if results_path.exists() else {
        m: rows for m, rows in historical.items() if m != args.model}
    mine = results.setdefault(args.model, {})
    jobs, inputs = [], {}
    for tag in tags:
        rec = records[tag]
        for order in ("forward", "reverse"):
            parts = _parts_for(rec, rdir, tag, LABELS,
                               LABELS if order == "forward" else list(reversed(LABELS)),
                               headers["default"], headers["bgfirst"])
            retained = manifest_parts(parts)
            identity = digest({"model": args.model, "effort": args.effort,
                               **({"provider": args.provider, "require_parameters": True}
                                  if args.provider == "openrouter" else {}),
                               "max_tokens": args.max_tokens, "system": system,
                               "schema": schema, "parts": retained})
            inputs[f"{tag}_{order}"] = {"request_identity": identity, "parts": retained}
            old = mine.get(tag, {}).get(order, {})
            if old.get("ok"):
                if old.get("request_identity") != identity:
                    raise ValueError(f"Input drift: {tag} {order}; use a new output directory")
                continue
            jobs.append((tag, order, parts, identity))
    if len(jobs) > args.max_calls:
        raise ValueError(f"{len(jobs)} requests exceed max-calls={args.max_calls}")
    manifest = {"model": args.model, "effort": args.effort, "api": "chat.completions",
                "provider": args.provider,
                "base_url": "https://openrouter.ai/api/v1" if args.provider == "openrouter" else "https://api.openai.com/v1",
                "max_completion_tokens": args.max_tokens,
                "request_options": {"reasoning": {"effort": args.effort}, "max_tokens": args.max_tokens,
                    "provider": {"require_parameters": True}} if args.provider == "openrouter" else {
                    "reasoning_effort": args.effort, "max_completion_tokens": args.max_tokens},
                "historical_max_tokens": 8000, "sdk_max_retries": 0,
                "project_id": args.project_id, "episode_id": args.episode_id,
                "baseline": str(args.baseline), "records": str(records_path),
                "shots": tags, "pending_calls": len(jobs), "workers": args.workers,
                "key_slot": args.key_slot, "archived_contract_sha256": {
                    name: digest(data) for name, data in snapshots.items()},
                "historical_source_sha256": original_hashes,
                "history_limit": "Historical image hashes were not recorded; retained paths and inline byte sizes are checked, and current exact request bytes are hashed.",
                "inputs": inputs}
    save(args.out / "input_manifest.json", manifest)
    print(json.dumps({"model": args.model, "effort": args.effort, "shots": tags,
                      "pending_calls": len(jobs), "workers": args.workers,
                      "live": args.live}, ensure_ascii=False), flush=True)
    if not args.live:
        return 0
    key_name = ("OPENROUTER_API_KEY" if args.provider == "openrouter" else
                "OPENAI_API_KEY_SECONDARY" if args.key_slot == "secondary" else "OPENAI_API_KEY")
    key = os.environ.get(key_name, "").strip()
    if not key:
        raise ValueError(f"{key_name} missing")
    started = time.monotonic()
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        pending = [pool.submit(ask, job, key, system, schema, args.max_tokens,
                               args.model, args.effort, args.provider) for job in jobs]
        for future in as_completed(pending):
            tag, order, slot = future.result()
            mine.setdefault(tag, {})[order] = slot
            save(results_path, results)
            print(json.dumps({"tag": tag, "order": order, "ok": slot["ok"],
                "winner": slot.get("normalized", {}).get("winner"),
                "elapsed_seconds": slot["elapsed_seconds"], "usage": slot.get("usage"),
                "error": slot.get("error")}, ensure_ascii=False), flush=True)
    summary = compare(results, records, sorted(model_tags[0]))
    summary["this_run_wall_seconds"] = round(time.monotonic() - started, 3)
    summary["this_run_call_count"] = len(jobs)
    summary["historical_sources_unchanged"] = all(
        digest(Path(p).read_bytes()) == old for p, old in original_hashes.items())
    save(args.out / "summary.json", summary)

    from build_astra_compare_gallery import build
    print(f"Gallery: {build(args.out)}", flush=True)
    print(json.dumps(summary, ensure_ascii=False, indent=2), flush=True)
    return 0 if all(mine.get(t, {}).get(o, {}).get("ok") for t in tags
                    for o in ("forward", "reverse")) else 1


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