import argparse
import json
import statistics
from time import perf_counter
from urllib.error import HTTPError
from urllib.request import Request, urlopen

BASE = "http://127.0.0.1:11434"
LABELS = ["RETURN", "RENEW", "RESERVE", "UNKNOWN"]
CASES = [
    ("Where can I return my book?", "RETURN"),
    ("I am still reading. Extend the due date.", "RENEW"),
    ("Book the 3D printer for Saturday.", "RESERVE"),
    ("What time does the library open?", "UNKNOWN"),
]
SCHEMA = {
    "type": "object",
    "properties": {"label": {"enum": LABELS}},
    "required": ["label"],
    "additionalProperties": False,
}
SYSTEM = (
    "Classify fictional library requests. RETURN means giving back a loan. "
    "RENEW means extending a loan. RESERVE means booking equipment. "
    "Use UNKNOWN for any other request. Return JSON with one field, label."
)


def request(endpoint, body=None):
    encoded = json.dumps(body).encode("utf-8") if body is not None else None
    req = Request(BASE + endpoint, data=encoded, headers={"Content-Type": "application/json"})
    with urlopen(req, timeout=120) as response:
        return json.load(response)


def model_digest(model):
    entry = next((item for item in request("/api/tags")["models"] if item["name"] == model), None)
    if entry is None:
        raise ValueError(f"Model not installed: {model}")
    return entry["digest"]


def evaluate(model):
    digest = model_digest(model)
    runtime = request("/api/version")["version"]
    results = []
    for case, (prompt, expected) in enumerate(CASES, start=1):
        start = perf_counter()
        outcome = request("/api/chat", {
            "model": model,
            "messages": [
                {"role": "system", "content": SYSTEM},
                {"role": "user", "content": prompt},
            ],
            "format": SCHEMA,
            "options": {"temperature": 0, "num_ctx": 2048, "num_predict": 64},
            "stream": False,
        })
        elapsed = round(perf_counter() - start, 3)
        try:
            item = json.loads(outcome["message"]["content"])
        except (KeyError, TypeError, ValueError):
            item = None
        shape_ok = isinstance(item, dict) and set(item) == {"label"} and item["label"] in LABELS
        actual = item["label"] if shape_ok else None
        results.append({
            "case": case, "expected": expected, "actual": actual,
            "passed": actual == expected, "seconds": elapsed,
        })

    if model_digest(model) != digest or request("/api/version")["version"] != runtime:
        raise ValueError("The model or runtime changed during the run. Run the check again.")
    passed = sum(row["passed"] for row in results)
    print(json.dumps({
        "runtime": runtime,
        "model": model,
        "digest": digest,
        "passed": passed,
        "total": len(results),
        "median_seconds": round(statistics.median(row["seconds"] for row in results), 3),
        "cases": results,
    }, indent=2))
    return passed == len(results)


parser = argparse.ArgumentParser(description="Run four local model checks and print a JSON report.")
parser.add_argument("model", help="Exact tag shown by ollama list")
args = parser.parse_args()
try:
    success = evaluate(args.model)
except HTTPError as error:
    raise SystemExit(f"Ollama returned HTTP {error.code}: {error.read().decode()}") from None
except (OSError, ValueError) as error:
    raise SystemExit(str(error)) from None
raise SystemExit(0 if success else 1)