#!/usr/bin/env python3
"""Run your task set against one model and record the answers with the settings that produced them.

Purpose: the measurement half of the personal evaluation harness. Sends every task in the task
    file to an OpenAI-compatible endpoint with fixed sampling settings, applies the cheap
    deterministic checks (required strings, forbidden strings, length), and writes a results
    file plus lab-notebook lines carrying the model, quantisation, settings and date, so that a
    run from today can be compared with a run from a month ago.
Platform: all (pure Python over HTTP; the server may be on any track or another machine)
Minimum memory: 8 GB on the machine running the model; this script needs very little
Assumes: Python 3.9 or later and an OpenAI-compatible endpoint at --base-url: llama-server
    from Part 6, the gateway from Part 9, or anything else that speaks the same API. The
    quantisation and engine version cannot be discovered reliably over that API, so they are
    arguments: fill them in or your results will be unreproducible.

Usage: python3 run-eval.py --base-url http://127.0.0.1:8080/v1 --model qwen3-4b \
           --quant Q4_K_M --tasks tasks-template.json --out results-qwen3-4b.json \
           --labbook labbook.md
       python3 run-eval.py --base-url http://127.0.0.1:4000/v1 --model qwen3-8b \
           --quant Q8_0 --api-key "$LAB_KEY" --tasks my-tasks.json --out results-qwen3-8b.json
"""

from __future__ import annotations

import argparse
import json
import platform
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Optional


def post_json(url: str, payload: dict, api_key: Optional[str], timeout: int) -> dict:
    body = json.dumps(payload).encode("utf-8")
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    request = urllib.request.Request(url, data=body, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:400]
        raise RuntimeError(f"{url} returned HTTP {exc.code}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"cannot reach {url}: {exc.reason}") from exc


def deterministic_checks(task: dict, answer: str) -> dict:
    """The scoring that needs no model. Cheap, exactly repeatable, and worth more than it looks."""
    lowered = answer.lower()
    required = [s for s in task.get("must_contain", []) if s.lower() not in lowered]
    forbidden = [s for s in task.get("must_not_contain", []) if s.lower() in lowered]
    words = len(answer.split())
    limit = task.get("max_words")
    return {
        "missing_required": required,
        "found_forbidden": forbidden,
        "words": words,
        "over_length": bool(limit and words > limit),
        "passed": not required and not forbidden and not (limit and words > limit),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--base-url", default="http://127.0.0.1:8080/v1")
    parser.add_argument("--api-key", default=None)
    parser.add_argument("--model", required=True, help="model name or alias the server answers to")
    parser.add_argument("--quant", required=True,
                        help="the quantisation actually loaded, e.g. Q4_K_M. Not discoverable "
                             "over the API, and a result without it is not reproducible.")
    parser.add_argument("--engine", default="llama.cpp", help="engine name, for the record")
    parser.add_argument("--engine-version", default="unknown",
                        help="engine version string, for the record")
    parser.add_argument("--tasks", default="tasks-template.json")
    parser.add_argument("--out", default="results.json")
    parser.add_argument("--system", default=None, help="system prompt sent with every task")
    parser.add_argument("--temperature", type=float, default=None,
                        help="overrides the temperature in the task file's settings block")
    parser.add_argument("--seed", type=int, default=None, help="overrides the task file's seed")
    parser.add_argument("--max-tokens", type=int, default=None)
    parser.add_argument("--only", default=None, help="run one category only")
    parser.add_argument("--timeout", type=int, default=300)
    parser.add_argument("--labbook", default=None)
    parser.add_argument("--notes", default="", help="one line about what this run is testing")
    args = parser.parse_args()

    spec = json.loads(Path(args.tasks).read_text(encoding="utf-8"))
    settings = dict(spec.get("settings", {}))
    settings.pop("comment", None)
    if args.temperature is not None:
        settings["temperature"] = args.temperature
    if args.seed is not None:
        settings["seed"] = args.seed
    if args.max_tokens is not None:
        settings["max_tokens"] = args.max_tokens

    tasks = [t for t in spec["tasks"] if args.only is None or t["category"] == args.only]
    if not tasks:
        sys.exit(f"no tasks in category {args.only!r}")

    endpoint = args.base_url.rstrip("/") + "/chat/completions"
    started = time.time()
    results = []
    passed = 0

    for task in tasks:
        messages = []
        if args.system:
            messages.append({"role": "system", "content": args.system})
        messages.append({"role": "user", "content": task["prompt"]})

        payload = {"model": args.model, "messages": messages}
        for key in ("temperature", "top_p", "seed", "max_tokens"):
            if key in settings:
                payload[key] = settings[key]

        began = time.time()
        body = post_json(endpoint, payload, args.api_key, args.timeout)
        elapsed = time.time() - began

        answer = (body["choices"][0]["message"]["content"] or "").strip()
        usage = body.get("usage", {})
        checks = deterministic_checks(task, answer)
        passed += int(checks["passed"])

        results.append({
            "id": task["id"],
            "category": task["category"],
            "prompt": task["prompt"],
            "reference": task["reference"],
            "rubric": task["rubric"],
            "answer": answer,
            "checks": checks,
            "seconds": round(elapsed, 2),
            "prompt_tokens": usage.get("prompt_tokens"),
            "completion_tokens": usage.get("completion_tokens"),
        })
        mark = "ok " if checks["passed"] else "BAD"
        print(f"  {mark} {task['id']} [{task['category']:11s}] {elapsed:6.2f}s  "
              f"{task['prompt'][:52].replace(chr(10), ' ')}")

    elapsed = time.time() - started
    run = {
        "lab": "part-10/lab-benchmark-models-on-your-own-tasks",
        "run_id": time.strftime("%Y%m%dT%H%M%S"),
        "task_set": spec.get("name", args.tasks),
        "task_set_version": spec.get("version"),
        "model": args.model,
        "quant": args.quant,
        "engine": args.engine,
        "engine_version": args.engine_version,
        "base_url": args.base_url,
        "settings": settings,
        "system_prompt": args.system,
        "host": platform.platform(),
        "tasks": len(results),
        "checks_passed": passed,
        "seconds": round(elapsed, 1),
        "date": time.strftime("%Y-%m-%d"),
        "notes": args.notes,
    }

    Path(args.out).write_text(json.dumps({"run": run, "results": results}, indent=2),
                              encoding="utf-8")
    print(f"\n{passed} of {len(results)} passed the deterministic checks "
          f"({elapsed:.1f} s). Answers written to {args.out}")
    print("Deterministic checks are a floor, not a score. Run judge.py next.")

    if args.labbook:
        with Path(args.labbook).open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(run) + "\n")
        print(f"recorded in {args.labbook}")


if __name__ == "__main__":
    main()