#!/usr/bin/env python3
"""Run a reproducible, local-only Ollama correctness/throughput check.

Python 3.10+; no third-party dependencies. No API keys or cloud endpoints.
Warning: Ollama cloud model tags can still route beyond localhost internally.
Verify downloaded local weights with "ollama show" and "ollama ps".
"""
from __future__ import annotations

import argparse
import csv
import json
import statistics
import sys
from pathlib import Path
from urllib.parse import urlsplit
from urllib.request import Request, urlopen
from urllib.error import URLError, HTTPError

DEFAULT_PROMPT = (
    "Reply with JSON only: give the sorted unique integers from "
    "[3,1,3,2] as a list under the key result."
)
DEFAULT_EXPECTED = {"result": [1, 2, 3]}
FIELDS = (
    "model", "trial", "correct", "prompt_tokens", "output_tokens",
    "decode_tok_s", "load_s", "total_s"
)


def local_url(value: str) -> str:
    parsed = urlsplit(value)
    if (parsed.scheme != "http" or parsed.hostname not in
            ("127.0.0.1", "localhost", "::1") or
            parsed.username is not None or parsed.password is not None or
            parsed.path not in ("", "/") or parsed.query or parsed.fragment):
        raise ValueError("Only an HTTP loopback Ollama server is supported")
    return value.rstrip("/") + "/api/generate"


def benchmark(
    models: list[str], *, runs: int = 3, context: int = 4096,
    max_output: int = 128, base_url: str = "http://127.0.0.1:11434",
    prompt: str = DEFAULT_PROMPT, expected: object = None,
    request_fn=urlopen,
) -> list[dict]:
    if not models or not 1 <= runs <= 30 or not 128 <= context <= 262144:
        raise ValueError("Need models, 1-30 trials and context 128-262144")
    if not 1 <= max_output <= 4096:
        raise ValueError("Output limit must be between 1 and 4096")
    for model in models:
        if not model.strip() or any(x.isspace() for x in model):
            raise ValueError("Use exact non-whitespace Ollama model tags")
        if ":cloud" in model.lower() or model.lower().endswith("-cloud"):
            raise ValueError("Cloud-tagged models are outside this local evaluation")
    endpoint = local_url(base_url)
    expected = DEFAULT_EXPECTED if expected is None else expected
    rows = []
    for model in models:
        for trial in range(1, runs + 1):
            data = {
                "model": model, "prompt": prompt, "stream": False,
                "format": "json", "keep_alive": "0",
                "options": {"num_ctx": context, "num_predict": max_output,
                            "temperature": 0},
            }
            request = Request(
                endpoint, data=json.dumps(data).encode("utf-8"),
                headers={"Content-Type": "application/json"}, method="POST",
            )
            try:
                with request_fn(request, timeout=240) as response:
                    result = json.load(response)
            except (URLError, HTTPError) as exc:
                raise RuntimeError(
                    f"{model}: Ollama request failed ({exc}). Check server, "
                    "downloaded model tag and available unified memory."
                ) from exc
            try:
                correct = json.loads(result.get("response", "{}")) == expected
            except (ValueError, TypeError):
                correct = False
            token_count = result.get("eval_count") or 0
            nanoseconds = result.get("eval_duration") or 0
            speed = token_count * 1e9 / nanoseconds if nanoseconds else None
            row = {
                "model": model, "trial": trial, "correct": correct,
                "prompt_tokens": result.get("prompt_eval_count"),
                "output_tokens": token_count,
                "decode_tok_s": round(speed, 2) if speed is not None else None,
                "load_s": round((result.get("load_duration") or 0) / 1e9, 2),
                "total_s": round((result.get("total_duration") or 0) / 1e9, 2),
            }
            rows.append(row)
    return rows


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--models", nargs="+", default=["gemma4:12b", "qwen3.8:27b"])
    parser.add_argument("--runs", type=int, default=3)
    parser.add_argument("--ctx", type=int, default=4096)
    parser.add_argument("--max-output", type=int, default=128)
    parser.add_argument("--url", default="http://127.0.0.1:11434")
    parser.add_argument("--csv", type=Path, help="Write results to this local CSV file")
    parser.add_argument("--prompt-file", type=Path, help="Optional local text prompt")
    parser.add_argument("--expected-file", type=Path, help="Optional expected JSON answer")
    args = parser.parse_args(argv)
    prompt = args.prompt_file.read_text(encoding="utf-8") if args.prompt_file else DEFAULT_PROMPT
    expected = (
        json.loads(args.expected_file.read_text(encoding="utf-8"))
        if args.expected_file else DEFAULT_EXPECTED
    )
    try:
        rows = benchmark(
            args.models, runs=args.runs, context=args.ctx,
            max_output=args.max_output, base_url=args.url,
            prompt=prompt, expected=expected,
        )
    except (ValueError, RuntimeError) as exc:
        print(f"Benchmark error: {exc}", file=sys.stderr)
        return 2
    for row in rows:
        print(json.dumps(row, ensure_ascii=False))
    for model in args.models:
        subset = [r for r in rows if r["model"] == model]
        speeds = [r["decode_tok_s"] for r in subset if r["decode_tok_s"] is not None]
        print(
            f"SUMMARY {model}: {sum(r['correct'] for r in subset)}/{len(subset)} "
            "correct, median decode tok/s "
            f"{statistics.median(speeds) if speeds else 'n/a'}"
        )
    if args.csv:
        with args.csv.open("w", encoding="utf-8", newline="") as handle:
            writer = csv.DictWriter(handle, fieldnames=FIELDS)
            writer.writeheader()
            writer.writerows(rows)
        print(f"Wrote {len(rows)} local CSV rows to {args.csv}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())