"""Agreement with the teacher on the held-out dev sets (synthetic and real).

The dev items carry the teacher's (Haiku 5.5, thinking) label as their last
message. A student is run on the same messages through Ollama, schema-forced,
thinking off, and compared field by field. This is agreement, not accuracy:
the hand-labelled testset.jsonl is the accuracy check.

    python3 eval_dev.py simplixity-router:2b-v1-q4_K_M [--num-gpu 99]
"""

import argparse
import json
import os
import re
import statistics
import sys
import time
import urllib.error
import urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
SCHEMA = json.load(open(os.path.join(HERE, "schema.json")))
SETS = {"syn": os.path.join(ROOT, "data", "valid.jsonl"), "real": os.path.join(ROOT, "data", "private", "valid_real.jsonl")}


def run(model, messages, num_gpu):
    body = {"model": model, "messages": messages, "stream": False, "format": SCHEMA, "think": False,
            "options": {"temperature": 0, "num_ctx": 8192, "num_predict": 500}, "keep_alive": "10m"}
    if num_gpu is not None:
        body["options"]["num_gpu"] = num_gpu
    req = urllib.request.Request("http://localhost:11434/api/chat", json.dumps(body).encode(),
                                 {"Content-Type": "application/json"})
    t = time.time()
    with urllib.request.urlopen(req, timeout=300) as r:
        out = json.load(r)
    return out["message"]["content"], time.time() - t


def words(queries):
    return set(re.findall(r"\w+", " ".join(queries).lower()))


def compare(student, teacher, settings):
    s = {"json": student is not None}
    if student is None:
        return s
    search = bool(student.get("needs_search")) or bool(student.get("wants_image"))
    s["search"] = search == teacher["needs_search"]
    if teacher["needs_search"] and search:
        for field in ("complexity", "answer_shape", "recency"):
            s[field] = str(student.get(field, "")).strip().lower() == teacher[field]
        s["image"] = bool(student.get("wants_image")) == teacher["wants_image"]
        if settings["auto_lenses"]:
            s["lenses"] = set(student.get("lenses") or []) == set(teacher["lenses"])
        a, b = words(student.get("search_queries") or []), words(teacher["search_queries"])
        s["query_overlap"] = len(a & b) / len(a | b) if a | b else 1.0
        s["route"] = s["search"] and s["complexity"] and s["answer_shape"]
    return s


def main():
    p = argparse.ArgumentParser()
    p.add_argument("model")
    p.add_argument("--num-gpu", type=int, default=None)
    p.add_argument("--tag", default="")
    args = p.parse_args()
    tag = args.tag or args.model.replace(":", "-").replace("/", "_")
    for name, path in SETS.items():
        rows = [json.loads(line) for line in open(path)]
        out_path = os.path.join(HERE if name == "syn" else os.path.join(ROOT, "data", "private"),
                                f"dev-{name}-{tag}.jsonl")
        results = []
        with open(out_path, "w") as f:
            for i, row in enumerate(rows):
                try:
                    raw, dt = run(args.model, row["messages"][:-1], args.num_gpu)
                except urllib.error.HTTPError as e:  # e.g. "token repeat limit reached": a failed call
                    raw, dt = "", 0.0
                    print(f"\n{row['id']}: {e.read().decode()[:200]}", file=sys.stderr)
                try:
                    student = json.loads(raw)
                except json.JSONDecodeError:
                    student = None
                teacher = json.loads(row["messages"][-1]["content"])
                r = {"id": row["id"], "score": compare(student, teacher, row["settings"]), "student": student,
                     "seconds": round(dt, 3)}
                results.append(r)
                f.write(json.dumps(r, ensure_ascii=False) + "\n")
                print(f"{name} {i + 1}/{len(rows)}", file=sys.stderr, end="\r")
        if name == "real":
            os.chmod(out_path, 0o600)
        print(f"\n== {name} dev (n={len(rows)}, median {statistics.median(r['seconds'] for r in results):.2f}s)")
        for field in ("json", "search", "route", "complexity", "answer_shape", "recency", "image", "lenses",
                      "query_overlap"):
            vals = [r["score"][field] for r in results if field in r["score"]]
            if vals:
                print(f"  {field:14} {100 * sum(vals) / len(vals):5.1f}%  (n={len(vals)})")


if __name__ == "__main__":
    main()