#!/usr/bin/env python3
"""Reproducible HTTP benchmark for Ollama and OpenAI-compatible runtimes."""

from __future__ import annotations

import argparse
import concurrent.futures
import csv
import json
import math
import platform
import statistics
import subprocess
import sys
import threading
import time
import urllib.error
import urllib.request
import uuid
from dataclasses import asdict, dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

PROMPTS = {
    "short": (
        "In exactly four concise bullet points, explain why database transactions are useful. "
        "Do not use a heading."
    ),
    "medium": (
        "You are reviewing a Laravel service that transfers money between two accounts. "
        "Explain the failure modes involving concurrency, partial writes, retries, and audit logs. "
        "Then provide a robust implementation plan. Keep the response technical and concise."
    ),
    "code": (
        "Write a Python 3.12 function named merge_intervals that accepts a list of integer pairs, "
        "validates each pair, merges overlapping or adjacent intervals, and returns sorted tuples. "
        "Include type hints, a short docstring, and five assert-based examples. Output only code."
    ),
}

CSV_FIELDS = (
    "run_id",
    "runtime",
    "model",
    "phase",
    "concurrency",
    "round",
    "slot",
    "request_id",
    "prompt_name",
    "prompt_chars",
    "max_tokens",
    "started_at",
    "ttft_s",
    "total_s",
    "prompt_tokens",
    "output_tokens",
    "prompt_tps",
    "decode_tps",
    "request_tps",
    "group_elapsed_s",
    "group_output_tps",
    "error",
)


@dataclass(frozen=True)
class RuntimeConfig:
    name: str
    kind: str
    base_url: str
    model: str


@dataclass
class Measurement:
    run_id: str
    runtime: str
    model: str
    phase: str
    concurrency: int
    round: int
    slot: int
    request_id: str
    prompt_name: str
    prompt_chars: int
    max_tokens: int
    started_at: str
    ttft_s: float | None = None
    total_s: float | None = None
    prompt_tokens: int | None = None
    output_tokens: int | None = None
    prompt_tps: float | None = None
    decode_tps: float | None = None
    request_tps: float | None = None
    group_elapsed_s: float | None = None
    group_output_tps: float | None = None
    error: str | None = None


def utc_now() -> str:
    return datetime.now(UTC).isoformat()


def percentile(values: list[float], fraction: float) -> float | None:
    if not values:
        return None
    ordered = sorted(values)
    position = (len(ordered) - 1) * fraction
    lower = math.floor(position)
    upper = math.ceil(position)
    if lower == upper:
        return ordered[lower]
    return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower)


def safe_version(command: list[str]) -> str | None:
    try:
        result = subprocess.run(command, capture_output=True, text=True, timeout=15, check=False)
    except (OSError, subprocess.SubprocessError):
        return None
    value = (result.stdout or result.stderr).strip()
    return value or None


def http_json(url: str, timeout: float = 15) -> dict[str, Any]:
    request = urllib.request.Request(url, headers={"Accept": "application/json"})
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return json.load(response)


def read_ndjson(response: Any, started: float) -> tuple[float | None, dict[str, Any], int]:
    first_token_at: float | None = None
    output_chars = 0
    terminal: dict[str, Any] = {}
    for raw_line in response:
        if not raw_line.strip():
            continue
        event = json.loads(raw_line)
        content = event.get("response", "")
        if content:
            output_chars += len(content)
            if first_token_at is None:
                first_token_at = time.perf_counter()
        if event.get("done"):
            terminal = event
    return (first_token_at - started if first_token_at else None, terminal, output_chars)


def read_sse(response: Any, started: float) -> tuple[float | None, dict[str, Any], int]:
    first_token_at: float | None = None
    output_chars = 0
    terminal: dict[str, Any] = {}
    for raw_line in response:
        line = raw_line.decode("utf-8").strip()
        if not line.startswith("data:"):
            continue
        payload = line[5:].strip()
        if payload == "[DONE]":
            break
        event = json.loads(payload)
        terminal = event
        choices = event.get("choices") or []
        if choices:
            text = choices[0].get("text") or choices[0].get("delta", {}).get("content") or ""
            if text:
                output_chars += len(text)
                if first_token_at is None:
                    first_token_at = time.perf_counter()
    return (first_token_at - started if first_token_at else None, terminal, output_chars)


def benchmark_request(
    runtime: RuntimeConfig,
    prompt_name: str,
    prompt: str,
    max_tokens: int,
    seed: int,
    timeout: float,
    run_id: str,
    phase: str,
    concurrency: int,
    round_number: int,
    slot: int,
) -> Measurement:
    measurement = Measurement(
        run_id=run_id,
        runtime=runtime.name,
        model=runtime.model,
        phase=phase,
        concurrency=concurrency,
        round=round_number,
        slot=slot,
        request_id=str(uuid.uuid4()),
        prompt_name=prompt_name,
        prompt_chars=len(prompt),
        max_tokens=max_tokens,
        started_at=utc_now(),
    )
    headers = {"Content-Type": "application/json", "Accept": "text/event-stream"}
    if runtime.kind == "ollama":
        url = f"{runtime.base_url.rstrip('/')}/api/generate"
        payload = {
            "model": runtime.model,
            "prompt": prompt,
            "stream": True,
            "raw": True,
            "keep_alive": "30m",
            "options": {
                "temperature": 0,
                "seed": seed,
                "num_predict": max_tokens,
                "num_ctx": 8192,
                "draft_num_predict": 0,
            },
        }
    else:
        url = f"{runtime.base_url.rstrip('/')}/v1/completions"
        payload = {
            "model": runtime.model,
            "prompt": prompt,
            "stream": True,
            "stream_options": {"include_usage": True},
            "temperature": 0,
            "seed": seed,
            "max_tokens": max_tokens,
        }

    started = time.perf_counter()
    try:
        request = urllib.request.Request(
            url,
            data=json.dumps(payload).encode("utf-8"),
            headers=headers,
            method="POST",
        )
        with urllib.request.urlopen(request, timeout=timeout) as response:
            if runtime.kind == "ollama":
                ttft, terminal, output_chars = read_ndjson(response, started)
                measurement.prompt_tokens = terminal.get("prompt_eval_count")
                measurement.output_tokens = terminal.get("eval_count")
                prompt_ns = terminal.get("prompt_eval_duration")
                decode_ns = terminal.get("eval_duration")
                if prompt_ns and measurement.prompt_tokens:
                    measurement.prompt_tps = measurement.prompt_tokens / (prompt_ns / 1e9)
                if decode_ns and measurement.output_tokens:
                    measurement.decode_tps = measurement.output_tokens / (decode_ns / 1e9)
            else:
                ttft, terminal, output_chars = read_sse(response, started)
                usage = terminal.get("usage") or {}
                measurement.prompt_tokens = usage.get("prompt_tokens")
                measurement.output_tokens = usage.get("completion_tokens")
            measurement.ttft_s = ttft
            measurement.total_s = time.perf_counter() - started
            if measurement.output_tokens and measurement.total_s:
                measurement.request_tps = measurement.output_tokens / measurement.total_s
                if measurement.decode_tps is None and measurement.ttft_s is not None:
                    decode_time = measurement.total_s - measurement.ttft_s
                    if decode_time > 0:
                        measurement.decode_tps = measurement.output_tokens / decode_time
            if output_chars == 0:
                measurement.error = "stream contained no generated text"
    except urllib.error.HTTPError as error:
        detail = error.read().decode("utf-8", errors="replace")[:1000]
        measurement.error = f"HTTP {error.code}: {detail}"
        measurement.total_s = time.perf_counter() - started
    except (OSError, ValueError, json.JSONDecodeError) as error:
        measurement.error = f"{type(error).__name__}: {error}"
        measurement.total_s = time.perf_counter() - started
    return measurement


def run_round(
    runtime: RuntimeConfig,
    prompt_items: list[tuple[str, str]],
    max_tokens: int,
    seed: int,
    timeout: float,
    run_id: str,
    phase: str,
    concurrency: int,
    round_number: int,
) -> tuple[list[Measurement], float]:
    barrier = threading.Barrier(concurrency)

    def execute(slot: int) -> Measurement:
        prompt_name, prompt = prompt_items[(round_number * concurrency + slot) % len(prompt_items)]
        barrier.wait()
        return benchmark_request(
            runtime,
            prompt_name,
            prompt,
            max_tokens,
            seed,
            timeout,
            run_id,
            phase,
            concurrency,
            round_number,
            slot,
        )

    group_started = time.perf_counter()
    with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as executor:
        measurements = list(executor.map(execute, range(concurrency)))
    group_elapsed = time.perf_counter() - group_started
    successful_tokens = sum(item.output_tokens or 0 for item in measurements if not item.error)
    group_tps = successful_tokens / group_elapsed if group_elapsed else 0.0
    for item in measurements:
        item.group_elapsed_s = group_elapsed
        item.group_output_tps = group_tps
    return measurements, group_elapsed


def summarize(measurements: list[Measurement]) -> list[dict[str, Any]]:
    groups: dict[tuple[str, int], list[Measurement]] = {}
    for measurement in measurements:
        if measurement.phase != "measured":
            continue
        groups.setdefault((measurement.runtime, measurement.concurrency), []).append(measurement)

    summaries: list[dict[str, Any]] = []
    for (runtime, concurrency), items in sorted(groups.items()):
        successful = [item for item in items if not item.error]
        summary: dict[str, Any] = {
            "runtime": runtime,
            "model": items[0].model,
            "concurrency": concurrency,
            "requests": len(items),
            "successes": len(successful),
            "failures": len(items) - len(successful),
        }
        for field in ("ttft_s", "total_s", "prompt_tps", "decode_tps", "request_tps"):
            values = [
                float(value)
                for item in successful
                if (value := getattr(item, field)) is not None
            ]
            summary[f"{field}_mean"] = statistics.fmean(values) if values else None
            summary[f"{field}_p50"] = percentile(values, 0.50)
            summary[f"{field}_p95"] = percentile(values, 0.95)
        group_tps = [
            float(item.group_output_tps)
            for item in successful
            if item.slot == 0 and item.group_output_tps is not None
        ]
        summary["aggregate_output_tps_mean"] = statistics.fmean(group_tps) if group_tps else None
        summary["aggregate_output_tps_p50"] = percentile(group_tps, 0.50)
        summary["aggregate_output_tps_p95"] = percentile(group_tps, 0.95)
        summaries.append(summary)
    return summaries


def write_csv(path: Path, measurements: list[Measurement]) -> None:
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=CSV_FIELDS, lineterminator="\n")
        writer.writeheader()
        for measurement in measurements:
            writer.writerow(asdict(measurement))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, default=Path(__file__).parent / "results")
    parser.add_argument("--runtime", choices=("all", "ollama", "vllm"), default="all")
    parser.add_argument("--ollama-url", default="http://127.0.0.1:11434")
    parser.add_argument("--ollama-model", default="qwen3.8:27b-mlx")
    parser.add_argument("--vllm-url", default="http://127.0.0.1:8001")
    parser.add_argument("--vllm-model", default="qwen3.8-27b-nvfp4")
    parser.add_argument("--warmups", type=int, default=2)
    parser.add_argument("--rounds", type=int, default=5)
    parser.add_argument("--concurrency", type=int, nargs="+", default=(1, 4))
    parser.add_argument("--max-tokens", type=int, default=128)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--timeout", type=float, default=600)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.warmups < 0 or args.rounds < 1 or any(value < 1 for value in args.concurrency):
        raise SystemExit("warmups must be >= 0; rounds and concurrency must be >= 1")

    runtimes = {
        "ollama": RuntimeConfig("ollama", "ollama", args.ollama_url, args.ollama_model),
        "vllm": RuntimeConfig("vllm-metal", "openai", args.vllm_url, args.vllm_model),
    }
    selected = list(runtimes.values()) if args.runtime == "all" else [runtimes[args.runtime]]
    for runtime in selected:
        health_url = (
            f"{runtime.base_url.rstrip('/')}/api/tags"
            if runtime.kind == "ollama"
            else f"{runtime.base_url.rstrip('/')}/v1/models"
        )
        try:
            http_json(health_url)
        except Exception as error:
            raise SystemExit(f"{runtime.name} is not ready at {health_url}: {error}") from error

    run_id = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
    output_dir = args.output_dir / run_id
    output_dir.mkdir(parents=True, exist_ok=False)
    metadata = {
        "run_id": run_id,
        "created_at": utc_now(),
        "host": {
            "platform": platform.platform(),
            "machine": platform.machine(),
            "python": platform.python_version(),
        },
        "arguments": vars(args) | {"output_dir": str(args.output_dir)},
        "runtimes": [asdict(runtime) for runtime in selected],
        "versions": {
            "ollama": safe_version(["ollama", "--version"]),
            "vllm": safe_version([str(Path.home() / ".venv-vllm-metal/bin/vllm"), "--version"]),
        },
        "model_metadata": {},
        "prompts": PROMPTS,
        "measurement_notes": {
            "ttft": (
                "Client wall time from request start until first non-empty generated text chunk."
            ),
            "total": "Client wall time from request start until stream completion.",
            "decode_tps": "Ollama server eval metric; otherwise output tokens / (total - TTFT).",
            "aggregate_output_tps": "Sum of successful output tokens / concurrent group wall time.",
        },
    }
    for runtime in selected:
        url = (
            f"{runtime.base_url.rstrip('/')}/api/show"
            if runtime.kind == "ollama"
            else f"{runtime.base_url.rstrip('/')}/v1/models"
        )
        try:
            if runtime.kind == "ollama":
                request = urllib.request.Request(
                    url,
                    data=json.dumps({"model": runtime.model}).encode(),
                    headers={"Content-Type": "application/json"},
                    method="POST",
                )
                with urllib.request.urlopen(request, timeout=30) as response:
                    metadata["model_metadata"][runtime.name] = json.load(response)
            else:
                metadata["model_metadata"][runtime.name] = http_json(url)
        except Exception as error:
            metadata["model_metadata"][runtime.name] = {"metadata_error": str(error)}

    prompt_items = list(PROMPTS.items())
    measurements: list[Measurement] = []
    total_groups = len(selected) * sum(args.warmups + args.rounds for _ in args.concurrency)
    completed_groups = 0
    for runtime in selected:
        for concurrency in args.concurrency:
            for phase, rounds in (("warmup", args.warmups), ("measured", args.rounds)):
                for round_number in range(rounds):
                    items, elapsed = run_round(
                        runtime,
                        prompt_items,
                        args.max_tokens,
                        args.seed,
                        args.timeout,
                        run_id,
                        phase,
                        concurrency,
                        round_number,
                    )
                    measurements.extend(items)
                    completed_groups += 1
                    errors = sum(bool(item.error) for item in items)
                    print(
                        f"[{completed_groups}/{total_groups}] {runtime.name} phase={phase} "
                        f"concurrency={concurrency} round={round_number + 1}/{rounds} "
                        f"elapsed={elapsed:.3f}s errors={errors}",
                        flush=True,
                    )

    summaries = summarize(measurements)
    metadata["completed_at"] = utc_now()
    metadata["summary"] = summaries
    (output_dir / "raw.json").write_text(
        json.dumps([asdict(item) | {
            "group_elapsed_s": getattr(item, "group_elapsed_s", None),
            "group_output_tps": getattr(item, "group_output_tps", None),
        } for item in measurements], indent=2),
        encoding="utf-8",
    )
    (output_dir / "metadata.json").write_text(
        json.dumps(metadata, indent=2, default=str), encoding="utf-8"
    )
    (output_dir / "summary.json").write_text(json.dumps(summaries, indent=2), encoding="utf-8")
    write_csv(output_dir / "raw.csv", measurements)
    print(json.dumps(summaries, indent=2))
    print(f"Results: {output_dir}")
    return 1 if any(item.error for item in measurements if item.phase == "measured") else 0


if __name__ == "__main__":
    sys.exit(main())