#!/usr/bin/env python3
"""Run a small, fixed-workload, end-to-end generation measurement.

This intentionally reports request-level output tokens per second. It does not
claim decode-only throughput or transfer results between runtimes.
"""
from __future__ import annotations

import json
import os
import statistics
import sys
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[2]
OUT = ROOT / "results" / "server_bench.json"

PROMPT = """Implement this Python 3 function and return code only, with no Markdown fence:

def merge_intervals(items: list[tuple[int, int]]) -> list[tuple[int, int]]:
    \"\"\"Merge overlapping or touching closed intervals.

    Normalize each reversed pair first, do not mutate the input, sort the
    result, return [] for empty input, and include a short docstring.
    \"\"\"
"""


def _positive_int(name: str, default: int) -> int:
    value = int(os.environ.get(name, str(default)))
    if value < 1:
        raise ValueError(f"{name} must be at least 1")
    return value


def _request(
    base: str,
    key: str | None,
    model: str,
    timeout: int,
    max_tokens: int,
) -> tuple[dict[str, Any], float]:
    headers = {"Content-Type": "application/json"}
    if key:
        headers["Authorization"] = f"Bearer {key}"
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": PROMPT}],
        "temperature": 0,
        "seed": 20260803,
        "max_tokens": max_tokens,
    }
    request = urllib.request.Request(
        base.rstrip("/") + "/chat/completions",
        data=json.dumps(payload).encode(),
        headers=headers,
        method="POST",
    )
    started = time.perf_counter()
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            body = json.loads(response.read().decode())
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode(errors="replace")[:2000]
        raise RuntimeError(f"HTTP {exc.code}: {detail}") from exc
    elapsed = time.perf_counter() - started
    return body, elapsed


def _sample(body: dict[str, Any], elapsed: float, index: int) -> dict[str, Any]:
    usage = body.get("usage") or {}
    prompt_tokens = usage.get("prompt_tokens")
    completion_tokens = usage.get("completion_tokens")
    if not isinstance(prompt_tokens, int) or prompt_tokens < 1:
        raise ValueError("response is missing a positive usage.prompt_tokens")
    if not isinstance(completion_tokens, int) or completion_tokens < 1:
        raise ValueError("response is missing a positive usage.completion_tokens")
    choices = body.get("choices") or []
    if len(choices) != 1:
        raise ValueError(f"expected one response choice, got {len(choices)}")
    message = choices[0].get("message") or {}
    content = message.get("content")
    if not isinstance(content, str) or not content.strip():
        raise ValueError("response content is empty")
    return {
        "index": index,
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "elapsed_seconds": elapsed,
        "output_tokens_per_second_end_to_end": completion_tokens / elapsed,
        "finish_reason": choices[0].get("finish_reason"),
        "response_text": content,
    }


def main() -> int:
    base = os.environ.get("OPENAI_BASE_URL", "http://127.0.0.1:8001/v1")
    key = os.environ.get("OPENAI_API_KEY")
    model = os.environ.get("SMOKE_MODEL", "local-qwen3-coder-next")
    timeout = _positive_int("REQUEST_TIMEOUT", 300)
    warmup_runs = _positive_int("WARMUP_RUNS", 1)
    measured_runs = _positive_int("MEASURED_RUNS", 5)
    max_tokens = _positive_int("MAX_TOKENS", 256)

    payload: dict[str, Any] = {
        "schema_version": 1,
        "suite": "fixed_generation_end_to_end",
        "status": "failed_run",
        "measured_at_utc": datetime.now(timezone.utc).isoformat(),
        "endpoint": base,
        "model": model,
        "settings": {
            "prompt": PROMPT,
            "temperature": 0,
            "seed": 20260803,
            "max_tokens": max_tokens,
            "concurrency": 1,
            "warmup_runs": warmup_runs,
            "measured_runs": measured_runs,
        },
        "warmup_completed": 0,
        "samples": [],
        "metric_definition": (
            "completion_tokens divided by full HTTP request elapsed seconds; "
            "includes prefill, decode, scheduling, serialization, and loopback overhead"
        ),
        "limitations": (
            "One short fixed coding prompt at concurrency one; not decode-only throughput, "
            "a coding-quality benchmark, a long-context test, or a sustained-load result."
        ),
    }

    try:
        for _ in range(warmup_runs):
            body, elapsed = _request(base, key, model, timeout, max_tokens)
            _sample(body, elapsed, 0)
            payload["warmup_completed"] += 1

        for index in range(1, measured_runs + 1):
            body, elapsed = _request(base, key, model, timeout, max_tokens)
            payload["samples"].append(_sample(body, elapsed, index))

        rates = [sample["output_tokens_per_second_end_to_end"] for sample in payload["samples"]]
        payload["summary"] = {
            "output_tokens_per_second_end_to_end_p50": statistics.median(rates),
            "output_tokens_per_second_end_to_end_min": min(rates),
            "output_tokens_per_second_end_to_end_max": max(rates),
        }
        payload["status"] = "measured"
    except Exception as exc:
        payload["error"] = str(exc)

    OUT.parent.mkdir(parents=True, exist_ok=True)
    OUT.write_text(json.dumps(payload, indent=2) + "\n")
    print(json.dumps(payload, indent=2))
    return 0 if payload["status"] == "measured" else 1


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