"""Run every task in tasks.json against two or more models and write the answers out blind.

Purpose: produce a scoring sheet in which every answer carries a random identifier and
         nothing else, so the reality check's twenty tasks can be marked without knowing
         which model wrote which answer. Local models are called through Ollama's HTTP
         API on this machine; answers produced elsewhere can be supplied as a JSON file.
Platform: all (spark, strix, mac, nvidia). Pure standard library: no packages to install.
Minimum memory: 12 GB, enough for a 4B-class and a 14B-class model one at a time. The
         30B-class comparison wants 24 GB or more.
Assumes: Ollama is reachable at $OLLAMA_HOST or http://localhost:11434, and every model
         named with --model has already been pulled with `ollama pull`. Writes three files
         into the output directory and overwrites them if they exist.

What it sends: one non-streaming POST /api/chat per task and model, with the prompt as a
single user message, `think` set to false unless --think is given, and an `options` object
that fixes the sampling settings so a rerun reproduces the run. What it keeps: the answer
text, the length of any thinking trace, `done_reason`, and Ollama's token and timing counters.

Usage: python run-blind-comparison.py --model qwen3:4b-q4_K_M --model qwen3:30b-a3b-q4_K_M
       python run-blind-comparison.py --model qwen3:4b-q4_K_M --model qwen3:14b-q4_K_M --dry-run
       python run-blind-comparison.py --model qwen3:4b-q4_K_M --model qwen3:14b-q4_K_M \\
           --external frontier=frontier-answers.json --out-dir .
"""
import argparse
import json
import os
import random
import re
import secrets
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

DEFAULT_HOST = os.environ.get("OLLAMA_HOST", "http://localhost:11434")
REQUIRED_TASK_FIELDS = ("id", "category", "prompt", "rubric")
COUNTERS = ("prompt_eval_count", "prompt_eval_duration", "eval_count", "eval_duration",
            "total_duration", "load_duration")

# Sampling settings. Non-thinking runs are greedy (temperature 0) so that a rerun gives the
# same answers. The Qwen3 model card says not to use greedy decoding in thinking mode and
# recommends temperature 0.6, top_p 0.95, top_k 20, min_p 0 there, so --think switches to
# those unless you set the values yourself.
GREEDY = {"temperature": 0.0}
THINKING_SAMPLING = {"temperature": 0.6, "top_p": 0.95, "top_k": 20, "min_p": 0.0}


def normalise_host(host: str) -> str:
    """Ollama's own OLLAMA_HOST is often set as host:port with no scheme."""
    host = host.strip().rstrip("/")
    if not host.startswith(("http://", "https://")):
        host = f"http://{host}"
    return host


def post_json(url: str, payload: dict, timeout: float) -> dict:
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return json.loads(resp.read().decode("utf-8"))


def get_json(url: str, timeout: float) -> dict:
    with urllib.request.urlopen(url, timeout=timeout) as resp:
        return json.loads(resp.read().decode("utf-8"))


def server_version(host: str, timeout: float) -> str:
    return str(get_json(f"{host}/api/version", timeout).get("version", "unknown"))


def installed_models(host: str, timeout: float) -> dict:
    """Model name -> {digest, size} for everything Ollama has locally, from GET /api/tags."""
    tags = get_json(f"{host}/api/tags", timeout)
    out = {}
    for m in tags.get("models", []):
        out[m.get("name", "")] = {"digest": str(m.get("digest", ""))[:12], "size": m.get("size")}
    return out


def resolve_model(name: str, present: dict):
    """The name as Ollama lists it: `qwen3:4b` is stored as itself, `qwen3` as `qwen3:latest`."""
    for candidate in (name, f"{name}:latest"):
        if candidate in present:
            return candidate
    return None


def strip_inline_thinking(text: str) -> tuple:
    """Remove a leading <think>...</think> block if an engine put the trace in the answer text."""
    match = re.match(r"^\s*<think>.*?</think>\s*", text, flags=re.DOTALL)
    if not match:
        return text, 0
    return text[match.end():], len(match.group(0))


def ask(host: str, model: str, prompt: str, options: dict, think, timeout: float) -> dict:
    """One non-streaming POST to /api/chat. Returns the answer text and the counters."""
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "stream": False,
        "options": options,
    }
    if think is not None:
        payload["think"] = think
    started = time.monotonic()
    try:
        body = post_json(f"{host}/api/chat", payload, timeout)
        think_sent = think
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        # A model whose template has no thinking support may reject the field; run it without.
        if exc.code == 400 and think is not None and "think" in detail.lower():
            payload.pop("think")
            body = post_json(f"{host}/api/chat", payload, timeout)
            think_sent = "omitted"
        else:
            raise urllib.error.URLError(f"HTTP {exc.code}: {detail[:200]}") from exc
    wall = time.monotonic() - started

    message = body.get("message", {})
    text, inline_chars = strip_inline_thinking(message.get("content", "") or "")
    thinking_chars = len(message.get("thinking", "") or "") + inline_chars
    result = {
        "text": text.strip(),
        "thinking_chars": thinking_chars,
        "done_reason": body.get("done_reason"),
        "think": think_sent,
        "wall_seconds": round(wall, 3),
    }
    for key in COUNTERS:
        result[key] = body.get(key)
    return result


def empty_result(text: str) -> dict:
    result = {"text": text, "thinking_chars": 0, "done_reason": None, "think": None,
              "wall_seconds": None}
    for key in COUNTERS:
        result[key] = None
    return result


def load_tasks(path: Path) -> dict:
    """Read the task file and refuse to run on one that is malformed."""
    try:
        doc = json.loads(path.read_text(encoding="utf-8"))
    except FileNotFoundError:
        raise SystemExit(f"{path} not found. Download tasks.json into this directory first.")
    except json.JSONDecodeError as exc:
        raise SystemExit(f"{path} is not valid JSON: {exc}")
    tasks = doc.get("tasks")
    if not isinstance(tasks, list) or not tasks:
        raise SystemExit(f"{path} has no tasks in it.")
    seen = set()
    for i, task in enumerate(tasks, start=1):
        missing = [f for f in REQUIRED_TASK_FIELDS if not task.get(f)]
        if missing:
            raise SystemExit(f"{path}: task {i} is missing {', '.join(missing)}.")
        if task["id"] in seen:
            raise SystemExit(f"{path}: task id {task['id']!r} appears twice.")
        seen.add(task["id"])
        if not all(f"{level} =" in task["rubric"] for level in ("0", "1", "2")):
            raise SystemExit(
                f"{path}: task {task['id']} has a rubric without explicit '2 =', '1 =' and '0 =' "
                "levels. Every rubric must say what each score looks like before any answer exists."
            )
    return doc


def describe_tasks(doc: dict) -> None:
    tasks = doc["tasks"]
    counts = {}
    for task in tasks:
        counts[task["category"]] = counts.get(task["category"], 0) + 1
    max_score = doc.get("scoring", {}).get("max_score", 2)
    margin = doc.get("scoring", {}).get("preregistered_margin")
    print(f"tasks: {len(tasks)} in {len(counts)} categories, {max_score} points each, "
          f"{len(tasks) * max_score} possible", file=sys.stderr)
    for category, n in counts.items():
        print(f"  {category:<22} {n:>2} tasks  {n * max_score:>2} points", file=sys.stderr)
    if margin is not None:
        print(f"pre-registered margin: {margin} points", file=sys.stderr)


def load_external(spec: str) -> tuple:
    """--external NAME=FILE, where FILE maps task id to answer text."""
    if "=" not in spec:
        raise SystemExit(f"--external wants NAME=FILE, got {spec!r}")
    name, _, filename = spec.partition("=")
    try:
        answers = json.loads(Path(filename).read_text(encoding="utf-8"))
    except FileNotFoundError:
        raise SystemExit(f"--external: {filename} not found.")
    except json.JSONDecodeError as exc:
        raise SystemExit(f"--external: {filename} is not valid JSON: {exc}")
    if not isinstance(answers, dict):
        raise SystemExit(f"{filename} must be a JSON object mapping task id to answer text.")
    return name.strip(), answers


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--tasks", default="tasks.json", help="task file to read")
    parser.add_argument("--model", action="append", default=[], metavar="NAME",
                        help="an Ollama model to run; repeat for each model compared")
    parser.add_argument("--external", action="append", default=[], metavar="NAME=FILE",
                        help="answers produced elsewhere, as a JSON object of task id to text")
    parser.add_argument("--host", default=DEFAULT_HOST, help="Ollama base URL")
    parser.add_argument("--out-dir", default=".", help="directory for the three output files")
    parser.add_argument("--think", action="store_true",
                        help="run every model in thinking mode (default: thinking off for all)")
    parser.add_argument("--temperature", type=float, default=None,
                        help="sampling temperature (default 0 without --think, 0.6 with it)")
    parser.add_argument("--top-p", type=float, default=None, help="nucleus sampling cut-off")
    parser.add_argument("--top-k", type=int, default=None, help="top-k sampling cut-off")
    parser.add_argument("--min-p", type=float, default=None, help="min-p sampling cut-off")
    parser.add_argument("--seed", type=int, default=1, help="sampling seed passed to Ollama")
    parser.add_argument("--num-predict", type=int, default=1024,
                        help="cap on generated tokens per answer, thinking included")
    parser.add_argument("--num-ctx", type=int, default=4096, help="context length to allocate")
    parser.add_argument("--timeout", type=float, default=900.0, help="seconds per request")
    parser.add_argument("--shuffle-seed", type=int, default=None,
                        help="seed for the blinding shuffle; omit for an unpredictable one")
    parser.add_argument("--dry-run", action="store_true",
                        help="check the task file, the server and the models, then stop")
    args = parser.parse_args()

    host = normalise_host(args.host)
    externals = [load_external(s) for s in args.external]
    if len(args.model) + len(externals) < 2:
        raise SystemExit("Give at least two answer sources: two --model values, or one of each.")

    doc = load_tasks(Path(args.tasks))
    tasks = doc["tasks"]
    describe_tasks(doc)

    sampling = dict(THINKING_SAMPLING if args.think else GREEDY)
    for key, value in (("temperature", args.temperature), ("top_p", args.top_p),
                       ("top_k", args.top_k), ("min_p", args.min_p)):
        if value is not None:
            sampling[key] = value
    options = {**sampling, "seed": args.seed, "num_predict": args.num_predict,
               "num_ctx": args.num_ctx}
    think = True if args.think else False

    version = None
    model_info = {}
    if args.model:
        try:
            version = server_version(host, timeout=30)
            present = installed_models(host, timeout=30)
        except (urllib.error.URLError, TimeoutError, OSError) as exc:
            raise SystemExit(
                f"Cannot reach Ollama at {host} ({exc}). Start it with `ollama serve`, "
                "or set --host if it runs elsewhere."
            ) from exc
        print(f"ollama server version {version} at {host}", file=sys.stderr)
        missing = []
        for name in args.model:
            resolved = resolve_model(name, present)
            if resolved is None:
                missing.append(name)
                continue
            info = present[resolved]
            model_info[name] = {"listed_as": resolved, **info}
            size_gb = (info["size"] or 0) / 1e9
            print(f"  {name:<28} digest {info['digest']}  {size_gb:5.1f} GB on disk",
                  file=sys.stderr)
        if missing:
            raise SystemExit(
                "These models are not present locally: " + ", ".join(missing)
                + "\nPull each one first, for example: ollama pull " + missing[0]
            )
    for name, mapping in externals:
        supplied = sum(1 for task in tasks if task["id"] in mapping)
        print(f"  {name:<28} external file, answers for {supplied}/{len(tasks)} tasks",
              file=sys.stderr)

    print(f"think: {think}   options: {json.dumps(options)}", file=sys.stderr)
    if args.dry_run:
        print("dry run: nothing was sent to a model.", file=sys.stderr)
        return

    answers = []
    for model in args.model:
        print(f"\n{model}", file=sys.stderr)
        for i, task in enumerate(tasks, start=1):
            print(f"  [{i:2d}/{len(tasks)}] {task['id']}", end="", file=sys.stderr, flush=True)
            try:
                result = ask(host, model, task["prompt"], options, think, args.timeout)
            except (urllib.error.URLError, TimeoutError, OSError) as exc:
                print(f"  FAILED: {exc}", file=sys.stderr)
                result = empty_result(f"[no answer: {exc}]")
            else:
                flag = ""
                if result["done_reason"] == "length":
                    flag = "  CUT OFF by num_predict"
                elif not result["text"]:
                    flag = "  EMPTY answer"
                print(f"  {result['wall_seconds']:6.1f} s  {result['eval_count'] or 0:5d} tok"
                      f"  {result['done_reason']}{flag}", file=sys.stderr)
            answers.append({"task": task["id"], "source": model, **result})

    for name, mapping in externals:
        for task in tasks:
            text = str(mapping.get(task["id"], "[no answer supplied]")).strip()
            answers.append({"task": task["id"], "source": name, **empty_result(text)})

    rng = random.Random(args.shuffle_seed) if args.shuffle_seed is not None else random.SystemRandom()
    for answer in answers:
        answer["response_id"] = secrets.token_hex(3)

    by_task = {task["id"]: [] for task in tasks}
    for answer in answers:
        by_task[answer["task"]].append(answer)
    for group in by_task.values():
        rng.shuffle(group)

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    blind_path = out_dir / "blind-responses.md"
    key_path = out_dir / "blind-key.json"
    scores_path = out_dir / "scores.csv"

    lines = [
        "# Blind responses",
        "",
        "Score every response below against its task's rubric, 0, 1 or 2, and write the score",
        "into scores.csv next to the matching response id. Do not open blind-key.json until you",
        "have finished scoring: it says which model wrote which answer.",
        "",
    ]
    for task in tasks:
        lines += [
            f"## {task['id']} - {task['category']}",
            "",
            "**Prompt**",
            "",
            "```text",
            task["prompt"].rstrip(),
            "```",
            "",
            f"**Rubric** {task['rubric']}",
            "",
        ]
        for answer in by_task[task["id"]]:
            lines += [
                f"### Response {answer['response_id']}",
                "",
                "```text",
                (answer["text"] or "[empty answer]").rstrip(),
                "```",
                "",
            ]
    blind_path.write_text("\n".join(lines) + "\n", encoding="utf-8")

    kept = ("task", "source", "thinking_chars", "done_reason", "think", "wall_seconds") + COUNTERS
    key_path.write_text(json.dumps({
        "generated": time.strftime("%Y-%m-%dT%H:%M:%S"),
        "host": host,
        "ollama_version": version,
        "tasks_file": str(args.tasks),
        "preregistered_margin": doc.get("scoring", {}).get("preregistered_margin"),
        "think": think,
        "options": options,
        "sources": args.model + [name for name, _ in externals],
        "models": model_info,
        "responses": {a["response_id"]: {k: a[k] for k in kept} for a in answers},
    }, indent=2) + "\n", encoding="utf-8")

    ordered = [a["response_id"] for task in tasks for a in by_task[task["id"]]]
    scores_path.write_text(
        "response_id,score\n" + "".join(f"{rid},\n" for rid in ordered), encoding="utf-8")

    cut = sum(1 for a in answers if a["done_reason"] == "length")
    empty = sum(1 for a in answers if not a["text"] or a["text"].startswith("[no answer"))
    print(f"\nwrote {blind_path}  ({len(answers)} responses)", file=sys.stderr)
    print(f"wrote {key_path}  (do not read it until you have scored)", file=sys.stderr)
    print(f"wrote {scores_path}  (fill in the score column)", file=sys.stderr)
    if cut:
        print(f"\n{cut} answer(s) were cut off by --num-predict {args.num_predict}. Score them "
              "as they are, or rerun every model with a higher cap so the comparison stays fair.",
              file=sys.stderr)
    if empty:
        print(f"{empty} answer(s) are empty or failed; see the troubleshooting section.",
              file=sys.stderr)


if __name__ == "__main__":
    main()