#!/usr/bin/env python3
# Published reference script.
"""Blindly judge every model pair in Series 2.

Three models create three pairwise judge calls per matter. Series 2 contains
60 matters, so a complete run produces 180 judge calls.

Examples:
    export OPENROUTER_API_KEY=...
    python3 judge.py --only q21
    python3 judge.py

The judge receives the evaluation prompt, one matter, and two answers labeled A
and B. Model names, reasoning traces, and sealed keys are not sent.
"""
from __future__ import annotations

import argparse
import csv
import hashlib
import itertools
import json
import os
import pathlib
import re
import time
from typing import Any

import requests


def position(question: str, model_1: str, model_2: str) -> tuple[str, str]:
    ordered = sorted((model_1, model_2))
    digest = hashlib.sha256((question + "|" + ordered[0] + "|" + ordered[1]).encode()).digest()
    return (ordered[0], ordered[1]) if digest[0] % 2 == 0 else (ordered[1], ordered[0])


def call_judge(url: str, key: str, model: str, prompt: str, timeout: int) -> tuple[dict[str, Any], list[dict[str, Any]]]:
    headers = {"Authorization": "Bearer " + key, "Content-Type": "application/json"}
    payload = {"model": model, "messages": [{"role": "user", "content": prompt}], "stream": False}
    attempts = []
    for attempt in range(1, 4):
        try:
            response = requests.post(url, headers=headers, json=payload, timeout=timeout)
            if response.status_code in (429, 500, 502, 503, 504):
                raise requests.HTTPError(f"HTTP {response.status_code}")
            response.raise_for_status()
            data = response.json()
            choice = data["choices"][0]
            attempts.append({"attempt": attempt, "transport_failed": False, "finish_reason": choice.get("finish_reason")})
            return {
                "report": (choice.get("message") or {}).get("content") or "",
                "judge_reasoning": (choice.get("message") or {}).get("reasoning") or "",
                "server_model": data.get("model"),
                "provider": data.get("provider"),
                "finish_reason": choice.get("finish_reason"),
                "usage": data.get("usage") or {},
                "attempt_number": attempt,
            }, attempts
        except Exception as exc:
            attempts.append({"attempt": attempt, "transport_failed": True, "error": str(exc)})
            if attempt == 3:
                raise
            time.sleep(10 * attempt)
    raise RuntimeError("unreachable")


def verdict_code(report: str) -> str:
    first_line = next((line.strip().upper() for line in report.splitlines() if line.strip()), "")
    if first_line.startswith("ANSWER A"):
        return "A"
    if first_line.startswith("ANSWER B"):
        return "B"
    if first_line.startswith("TIE"):
        return "TIE"
    if "BOTH" in first_line and "INADEQUATE" in first_line:
        return "BOTH_INADEQUATE"
    raise ValueError(f"Unparseable verdict: {first_line}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--judge-model", default="openai/gpt-5.6-sol")
    parser.add_argument("--base-url", default="https://openrouter.ai/api/v1")
    parser.add_argument("--questions", default="questions/")
    parser.add_argument("--answers", default="answers/")
    parser.add_argument("--eval-prompt", default="eval_prompt.txt")
    parser.add_argument("--verdicts", default="verdicts/")
    parser.add_argument("--only")
    parser.add_argument("--timeout", type=int, default=3600)
    args = parser.parse_args()

    key = os.environ.get("OPENROUTER_API_KEY", "").strip()
    if not key:
        raise SystemExit("OPENROUTER_API_KEY is not set")
    base = pathlib.Path(__file__).resolve().parent
    qdir = base / args.questions
    adir = base / args.answers
    vdir = base / args.verdicts
    vdir.mkdir(parents=True, exist_ok=True)
    eval_text = (base / args.eval_prompt).read_text(encoding="utf-8")
    eval_hash = hashlib.sha256(eval_text.encode()).hexdigest()
    questions = sorted(qdir.glob("q*.md"))
    if args.only:
        questions = [p for p in questions if p.stem == args.only]

    model_dirs = sorted([p for p in adir.iterdir() if p.is_dir()])
    if len(model_dirs) != 3:
        raise SystemExit(f"Expected three model directories, found {len(model_dirs)}")
    rows = []
    for question_path in questions:
        qid = question_path.stem
        matter = question_path.read_text(encoding="utf-8")
        for left, right in itertools.combinations(model_dirs, 2):
            model_a, model_b = position(qid, left.name, right.name)
            answer_a_path = next((adir / model_a).glob(f"*__{qid}_r1.md"))
            answer_b_path = next((adir / model_b).glob(f"*__{qid}_r1.md"))
            output = vdir / f"{qid}__{left.name}__vs__{right.name}.json"
            if output.exists():
                continue
            prompt = f"{eval_text}\n\n# ORIGINAL QUESTION\n{matter}\n\n# ANSWER A\n{answer_a_path.read_text(encoding='utf-8')}\n\n# ANSWER B\n{answer_b_path.read_text(encoding='utf-8')}"
            result, attempts = call_judge(args.base_url.rstrip("/") + "/chat/completions", key, args.judge_model, prompt, args.timeout)
            code = verdict_code(result["report"])
            winner = "TIE" if code == "TIE" else ("BOTH_INADEQUATE" if code == "BOTH_INADEQUATE" else (model_a if code == "A" else model_b))
            record = {
                "question": qid,
                "model_1": left.name,
                "model_2": right.name,
                "answer_a_was": model_a,
                "answer_b_was": model_b,
                "verdict_code": code,
                "winner": winner,
                "judge_model": args.judge_model,
                "eval_prompt_sha256": eval_hash,
                "attempt_history": attempts,
                **result,
            }
            output.write_text(json.dumps(record, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
            rows.append(record)
            print(qid, left.name, "vs", right.name, code)


if __name__ == "__main__":
    main()