#!/usr/bin/env python3
"""Run the same uncached Orca HTTP workloads against CIRU and Gufo."""

import argparse
import json
import statistics
import time
import urllib.error
import urllib.request


CASES = {
    "mixed-short": ("mixed", 256),
    "mixed-medium": ("mixed", 2048),
    "mixed-long": ("mixed", 8192),
    "mixed-very-long": ("mixed", 32768),
    "repetitive-short": ("repetitive", 256),
}


def make_prompt(style: str, words: int, nonce: int) -> str:
    seed = (
        "Local inference performance depends on prompt length, memory traffic, "
        "attention state, expert routing, speculative acceptance, and scheduling. "
    )
    prefix = (seed * (words // len(seed.split()) + 2))[: words * 8]
    if style == "mixed":
        instruction = (
            "Write 20 distinct numbered observations about designing a reliable "
            "local inference benchmark. Continue with substantial detail."
        )
    else:
        instruction = (
            "Repeat 'The quick brown fox jumps over the lazy dog.' on separate "
            "lines until stopped."
        )
    return f"Trial {nonce:06d}.\n" + prefix + "\n\n" + instruction


def request(base_url: str, model: str, style: str, words: int, nonce: int,
            max_tokens: int, stream: bool) -> dict:
    body = {
        "model": model,
        "messages": [{"role": "user", "content": make_prompt(style, words, nonce)}],
        "temperature": 0,
        "presence_penalty": 0,
        "max_tokens": max_tokens,
        "reasoning_effort": "off",
        "chat_template_kwargs": {"enable_thinking": False},
    }
    if stream:
        body["stream"] = True
        body["stream_options"] = {"include_usage": True}
    req = urllib.request.Request(
        base_url.rstrip("/") + "/v1/chat/completions",
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"},
    )
    start = time.monotonic()
    try:
        with urllib.request.urlopen(req, timeout=900) as response:
            server_timing = response.headers.get("Server-Timing")
            if stream:
                usage, timings, finish_reason, ttft_s = None, None, None, None
                content_parts = []
                for raw in response:
                    line = raw.decode(errors="replace").strip()
                    if not line.startswith("data:"):
                        continue
                    payload = line[5:].strip()
                    if payload == "[DONE]":
                        break
                    event = json.loads(payload)
                    if event.get("usage") is not None:
                        usage = event["usage"]
                    if event.get("timings") is not None:
                        timings = event["timings"]
                    for choice in event.get("choices", []):
                        delta = choice.get("delta") or {}
                        if ttft_s is None and delta.get("content"):
                            ttft_s = round(time.monotonic() - start, 3)
                        if delta.get("content"):
                            content_parts.append(delta["content"])
                        if choice.get("finish_reason"):
                            finish_reason = choice["finish_reason"]
                response_text = "".join(content_parts)
            else:
                result = json.load(response)
                usage = result.get("usage")
                timings = result.get("timings")
                finish_reason = result.get("choices", [{}])[0].get("finish_reason")
                ttft_s = None
                response_text = result.get("choices", [{}])[0].get("message", {}).get("content", "")
    except urllib.error.HTTPError as exc:
        raise RuntimeError(f"HTTP {exc.code}: {exc.read(500).decode(errors='replace')}") from exc
    row = {
        "style": style,
        "words": words,
        "nonce": nonce,
        "wall_s": round(time.monotonic() - start, 3),
        "ttft_s": ttft_s,
        "finish_reason": finish_reason,
        "response_text": response_text,
        "usage": usage,
        "timings": timings,
        "server_timing": server_timing,
    }
    return row


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url", required=True)
    parser.add_argument("--model", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--cases", default="mixed-short,mixed-medium,mixed-long,repetitive-short")
    parser.add_argument("--repeats", type=int, default=3)
    parser.add_argument("--max-tokens", type=int, default=128)
    parser.add_argument("--stream", action="store_true", help="measure time to first visible token")
    args = parser.parse_args()
    selected = args.cases.split(",")
    unknown = set(selected) - set(CASES)
    if unknown:
        parser.error(f"unknown cases: {sorted(unknown)}")

    rows = []
    nonce = 0
    for case in selected:
        style, words = CASES[case]
        nonce += 1
        request(args.base_url, args.model, style, words, nonce, args.max_tokens, args.stream)
        for _ in range(args.repeats):
            nonce += 1
            row = request(args.base_url, args.model, style, words, nonce, args.max_tokens, args.stream)
            rows.append({"case": case, **row})
            timings = row["timings"] or {}
            print(
                f"{case}: wall={row['wall_s']}s "
                f"ttft={row['ttft_s']}s "
                f"prompt={timings.get('prompt_per_second')} tok/s "
                f"decode={timings.get('predicted_per_second')} tok/s "
                f"cache={timings.get('cache_n')} tokens",
                flush=True,
            )

    with open(args.output, "w", encoding="utf-8") as output:
        json.dump({"model": args.model, "base_url": args.base_url,
                   "stream": args.stream,
                   "cases": selected, "repeats": args.repeats, "rows": rows}, output, indent=2)
    for case in selected:
        values = [row["wall_s"] for row in rows if row["case"] == case]
        print(f"{case}: median wall={statistics.median(values):.3f}s")


if __name__ == "__main__":
    main()