#!/usr/bin/env python3
"""Measure warm streaming TTFT and end-to-end latency on the pinned server."""

from __future__ import annotations

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

from harness_lib import PORT, SEEDS, Server, dump


ROOT = Path("/workspace/kingy-qwen38")


def prompt_for_tokens(server: Server, target: int = 512) -> tuple[str, int]:
    unit = "Cobalt maple quartz delta: write a concise synthetic audit note grounded only in this ledger. "
    repeats = max(1, target // max(1, server.tokenize(unit)))
    text = unit * repeats
    count = server.tokenize(text)
    for _ in range(4):
        if abs(count - target) <= 8:
            break
        repeats = max(1, round(repeats * target / max(1, count)))
        text = unit * repeats
        count = server.tokenize(text)
    return text, count


def stream_once(model_id: str, prompt: str) -> dict:
    payload = {
        "model": model_id,
        "messages": [
            {"role": "system", "content": "Answer directly in concise prose."},
            {"role": "user", "content": prompt},
        ],
        "seed": SEEDS[0],
        "temperature": 0,
        "top_p": 0.95,
        "max_tokens": 64,
        "stream": True,
    }
    req = urllib.request.Request(
        f"http://127.0.0.1:{PORT}/v1/chat/completions",
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    began = time.perf_counter()
    first = None
    chunks = 0
    content: list[str] = []
    with urllib.request.urlopen(req, timeout=900) as response:
        for raw in response:
            line = raw.decode(errors="replace").strip()
            if not line.startswith("data: ") or line == "data: [DONE]":
                continue
            event = json.loads(line[6:])
            delta = ((event.get("choices") or [{}])[0].get("delta") or {})
            token = delta.get("content") or ""
            if token:
                if first is None:
                    first = time.perf_counter()
                chunks += 1
                content.append(token)
    ended = time.perf_counter()
    if first is None:
        raise RuntimeError("stream returned no content token")
    return {
        "ttft_seconds": first - began,
        "end_to_end_seconds": ended - began,
        "content_chunks": chunks,
        "output": "".join(content),
    }


def power(path: Path) -> dict:
    values = []
    for line in path.read_text(errors="replace").splitlines() if path.exists() else []:
        bits = [value.strip() for value in line.split(",")]
        if len(bits) >= 3:
            try:
                values.append(float(bits[2]))
            except ValueError:
                pass
    return {
        "samples": len(values),
        "median_watts": statistics.median(values) if values else None,
        "peak_watts": max(values) if values else None,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-id", required=True)
    parser.add_argument("--model-path", type=Path, required=True)
    parser.add_argument("--root", type=Path, default=ROOT)
    parser.add_argument("--server-bin", type=Path, default=ROOT / "llama.cpp/build/bin/llama-server")
    args = parser.parse_args()
    out = args.root / "results" / args.model_id / "stream-timing"
    server = Server(args.model_id, args.model_path, args.server_bin, out, ctx=8192, reasoning="off", label="stream-timing")
    try:
        server.start()
        prompt, tokens = prompt_for_tokens(server)
        warmup = stream_once(args.model_id, prompt)
        runs = [stream_once(args.model_id, prompt) for _ in range(5)]
    finally:
        server_result = server.stop()
    result = {
        "model_id": args.model_id,
        "prompt_tokens_tokenize": tokens,
        "completion_cap": 64,
        "warmup_excluded": warmup,
        "runs": runs,
        "median_ttft_seconds": statistics.median(row["ttft_seconds"] for row in runs),
        "median_end_to_end_seconds": statistics.median(row["end_to_end_seconds"] for row in runs),
        "min_ttft_seconds": min(row["ttft_seconds"] for row in runs),
        "max_ttft_seconds": max(row["ttft_seconds"] for row in runs),
        "server": server_result,
        "power": power(server.gpu_path),
    }
    dump(out / "summary.json", result)
    print(json.dumps({key: result[key] for key in ("model_id", "median_ttft_seconds", "median_end_to_end_seconds")}))


if __name__ == "__main__":
    main()
