#!/usr/bin/env python3
"""Small warm-response timing workload; not a first-token or quality benchmark."""
import argparse
import json
from pathlib import Path
import statistics
import time
import urllib.error
import urllib.request

PROMPTS = [
    "Explain what a reverse proxy does in three sentences.",
    "Extract only the opening time as JSON: The workshop opens at 09:30 on Saturday.",
    "Return a Python function double(x) that returns twice its input. No commentary.",
]


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def run(model, prompt):
    request = urllib.request.Request("http://127.0.0.1:11434/api/generate", data=json.dumps({
        "model": model, "prompt": prompt, "stream": False,
        "options": {"temperature": 0, "num_predict": 128, "num_ctx": 4096},
    }).encode(), headers={"Content-Type": "application/json"})
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), NoRedirect())
    start = time.monotonic()
    with opener.open(request, timeout=180) as response:
        payload = json.loads(response.read(2 * 1024 * 1024))
    duration = time.monotonic() - start
    if not payload.get("response"):
        raise ValueError("No generated response returned")
    count, nanos = payload.get("eval_count", 0), payload.get("eval_duration", 0)
    return {"elapsed_seconds": duration, "generation_tokens_per_second": count / (nanos / 1e9) if nanos else None,
            "load_duration_ns": payload.get("load_duration"), "prompt_eval_duration_ns": payload.get("prompt_eval_duration"),
            "response": payload["response"]}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("models", nargs="+")
    parser.add_argument("--runs", type=int, default=3)
    parser.add_argument("--output", default="comparison.json")
    args = parser.parse_args()
    if not 1 <= args.runs <= 20 or len(args.models) > 5:
        parser.error("Use 1-20 runs and at most 5 models")
    target = Path(args.output)
    if target.exists():
        parser.error("Output already exists; choose a new report filename")
    report = {"note": "Warm non-streaming timing; not TTFT or automatic quality scoring", "models": []}
    failed = False
    for model in args.models:
        item = {"model": model, "runs": [], "errors": []}
        report["models"].append(item)
        try:
            run(model, "Reply with OK.")
        except (OSError, ValueError, urllib.error.URLError) as error:
            item["warmup_error"] = str(error)
            failed = True
            continue
        for repetition in range(args.runs):
            for prompt in PROMPTS:
                try:
                    item["runs"].append({"repetition": repetition + 1, "prompt": prompt, **run(model, prompt)})
                except (OSError, ValueError, urllib.error.URLError) as error:
                    item["errors"].append({"repetition": repetition + 1, "prompt": prompt, "error": str(error)})
                    failed = True
        if item["runs"]:
            item["median_seconds"] = statistics.median(result["elapsed_seconds"] for result in item["runs"])
    with target.open("x", encoding="utf-8") as handle:
        json.dump(report, handle, indent=2)
    print(f"Saved {target}; review answers and failures manually.")
    return 1 if failed else 0


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