#!/usr/bin/env python3
"""Replay a small, fixed workload against llama.cpp's native /completion API.

Use the SAME script with the baseline server, then with the draft server.
Results describe the server you run, not a LeCompute reference benchmark.
No third-party packages, telemetry, or model downloads. Outputs stay in JSON.
"""

import argparse
import datetime
import json
import statistics
import time
import urllib.error
import urllib.request
from pathlib import Path


PROMPTS = {
    "code": (
        "Rewrite this Python function with type annotations and a docstring. "
        "Preserve its behavior. Return only the function.\n"
        "def unique_names(names):\n"
        "    result = []\n"
        "    for name in names:\n"
        "        if name not in result:\n"
        "            result.append(name)\n"
        "    return result\n"
    ),
    "explanation": (
        "Explain how a CPU cache differs from main memory. "
        "Use one concrete programming example and at most 150 words."
    ),
    "open_text": (
        "Write a short scene about two engineers investigating an intermittent "
        "failure in a remote observatory. Use at most 150 words."
    ),
}


def measure(url, prompt, timeout=300):
    payload = {
        "prompt": prompt,
        "n_predict": 256,
        "temperature": 0,
        "seed": 42,
        "cache_prompt": False,
        "stream": False,
    }
    request = urllib.request.Request(
        url.rstrip("/") + "/completion",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
    )
    start = time.perf_counter()
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            result = json.load(response)
        elapsed = (time.perf_counter() - start) * 1000
        if not isinstance(result, dict) or not isinstance(result.get("content"), str):
            raise ValueError("Expected a native llama.cpp completion with string content")
        if result.get("error"):
            raise ValueError(str(result["error"]))
        return {
            "ok": True,
            "elapsed_ms": elapsed,
            "content": result["content"],
            "timings": result.get("timings"),
            "tokens_predicted": result.get("tokens_predicted"),
            "truncated": result.get("truncated"),
            "stopped_limit": result.get("stopped_limit"),
            "stop_type": result.get("stop_type"),
            "generation_settings": result.get("generation_settings"),
        }
    except (OSError, ValueError) as error:
        return {
            "ok": False,
            "elapsed_ms": (time.perf_counter() - start) * 1000,
            "error": str(error),
        }


def summarize(rows):
    summary = {}
    for task in dict.fromkeys(row["task"] for row in rows):
        task_rows = [row for row in rows if row["task"] == task]
        times = [row["elapsed_ms"] for row in task_rows if row["ok"]]
        summary[task] = {
            "successful": len(times),
            "failed": len(task_rows) - len(times),
            "median_elapsed_ms": statistics.median(times) if times else None,
            "min_elapsed_ms": min(times) if times else None,
            "max_elapsed_ms": max(times) if times else None,
        }
    return summary


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--url", default="http://127.0.0.1:8080")
    parser.add_argument("--label", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--repeats", type=int, default=5)
    args = parser.parse_args()
    if args.repeats < 1:
        parser.error("--repeats must be positive")
    # Exclusive creation prevents silently replacing a previous experiment.
    try:
        output = args.output.open("x", encoding="utf-8")
    except OSError as error:
        parser.error(str(error))
    report = {
        "label": args.label,
        "started_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "url": args.url,
        "settings": {"temperature": 0, "seed": 42, "n_predict": 256,
                     "cache_prompt": False, "stream": False},
        "prompts": PROMPTS,
        "warmup": [],
        "runs": [],
        "limitations": "Raw completion prompts; no chat template. No streaming TTFT, "
                        "GPU-memory or power measurement. Inspect output length, "
                        "truncation and task success before comparing times.",
    }
    try:
        for task, prompt in PROMPTS.items():
            row = {"task": task, **measure(args.url, prompt)}
            report["warmup"].append(row)
            if not row["ok"]:
                print("Warmup failed: " + row["error"])
                return 1
        for repeat in range(args.repeats):
            for task, prompt in PROMPTS.items():
                row = {"task": task, "repeat": repeat + 1, **measure(args.url, prompt)}
                report["runs"].append(row)
                print(f"{args.label}: {task} {repeat + 1}/{args.repeats}: "
                      f"{row['elapsed_ms']:.1f} ms, ok={row['ok']}")
        return 1 if any(not row["ok"] for row in report["runs"]) else 0
    finally:
        report["summary"] = summarize(report["runs"])
        with output:
            json.dump(report, output, ensure_ascii=False, indent=2)
            output.write("\n")
        print(f"Results saved to {args.output}")


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