#!/usr/bin/env python3
"""Small, dependency-free benchmark for an OpenAI-compatible chat endpoint.

This measures endpoint latency and reported completion usage. It does not grade
model quality, coding, reasoning, or agent behavior.
"""

from __future__ import annotations

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


DEFAULT_PROMPT = (
    "Explain why a quality-gated heterogeneous inference runtime should reject "
    "a faster quantized path when it changes a near-tie top-1 token."
)


def percentile(values: list[float], fraction: float) -> float:
    ordered = sorted(values)
    if not ordered:
        raise ValueError("cannot calculate a percentile of an empty list")
    position = (len(ordered) - 1) * fraction
    lower = int(position)
    upper = min(lower + 1, len(ordered) - 1)
    weight = position - lower
    return ordered[lower] * (1.0 - weight) + ordered[upper] * weight


def post_json(url: str, payload: dict[str, Any], timeout: float) -> tuple[dict[str, Any], float]:
    request = urllib.request.Request(
        url,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    started = time.perf_counter()
    with urllib.request.urlopen(request, timeout=timeout) as response:
        body = response.read()
    elapsed = time.perf_counter() - started
    return json.loads(body), elapsed


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-url", default="http://127.0.0.1:8080/v1")
    parser.add_argument("--model", default="qwen3.6-35b-a3b-umox")
    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
    parser.add_argument("--system", default="You are a precise engineering assistant.")
    parser.add_argument("--iterations", type=int, default=5)
    parser.add_argument("--warmup", type=int, default=1)
    parser.add_argument("--max-tokens", type=int, default=256)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--timeout", type=float, default=600.0)
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()

    if args.iterations < 1 or args.warmup < 0 or args.max_tokens < 1:
        parser.error("iterations/max-tokens must be positive and warmup cannot be negative")

    endpoint = args.base_url.rstrip("/") + "/chat/completions"
    payload = {
        "model": args.model,
        "messages": [
            {"role": "system", "content": args.system},
            {"role": "user", "content": args.prompt},
        ],
        "temperature": args.temperature,
        "seed": args.seed,
        "max_tokens": args.max_tokens,
        "stream": False,
    }

    records: list[dict[str, Any]] = []
    try:
        for index in range(args.warmup + args.iterations):
            response, elapsed = post_json(endpoint, payload, args.timeout)
            usage = response.get("usage") or {}
            timings = response.get("timings") or {}
            completion_tokens = usage.get("completion_tokens")
            measured_tps = (
                completion_tokens / elapsed
                if isinstance(completion_tokens, int) and completion_tokens > 0
                else None
            )
            record = {
                "iteration": index - args.warmup,
                "warmup": index < args.warmup,
                "wall_seconds": elapsed,
                "completion_tokens": completion_tokens,
                "wall_completion_tokens_per_second": measured_tps,
                "server_predicted_tokens_per_second": timings.get("predicted_per_second"),
                "finish_reason": ((response.get("choices") or [{}])[0]).get("finish_reason"),
            }
            if index >= args.warmup:
                records.append(record)
            print(json.dumps(record, sort_keys=True))
    except (urllib.error.URLError, TimeoutError, json.JSONDecodeError) as error:
        raise SystemExit(f"benchmark request failed: {error}") from error

    wall_times = [float(row["wall_seconds"]) for row in records]
    wall_tps = [
        float(row["wall_completion_tokens_per_second"])
        for row in records
        if row["wall_completion_tokens_per_second"] is not None
    ]
    summary: dict[str, Any] = {
        "schema_version": 1,
        "endpoint": endpoint,
        "model": args.model,
        "configuration": {
            "iterations": args.iterations,
            "warmup": args.warmup,
            "max_tokens": args.max_tokens,
            "temperature": args.temperature,
            "seed": args.seed,
        },
        "latency_seconds": {
            "median": statistics.median(wall_times),
            "p95": percentile(wall_times, 0.95),
            "minimum": min(wall_times),
            "maximum": max(wall_times),
        },
        "wall_completion_tokens_per_second": (
            {
                "median": statistics.median(wall_tps),
                "p95": percentile(wall_tps, 0.95),
                "minimum": min(wall_tps),
                "maximum": max(wall_tps),
            }
            if wall_tps
            else None
        ),
        "samples": records,
        "warning": "Endpoint performance only; this result does not establish model quality.",
    }

    rendered = json.dumps(summary, indent=2, sort_keys=True) + "\n"
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(rendered, encoding="utf-8")
        print(f"wrote {args.output}")
    else:
        print(rendered)
    return 0


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