"""race.py — the same prompts, one request at a time, against a denv endpoint
and against named providers through OpenRouter (each pinned, no fallback).
Prints median and p90 decode speed and time to first token per target.

  DENV_BASE_URL, DENV_API_KEY      the denv endpoint (key never printed)
  OPENROUTER_API_KEY               for the providers
  python race.py --denv-model qwen3.8-27b-fp8 --providers cerebras,modelrun,modal
"""
import argparse
import json
import os
import statistics
import time

from openai import OpenAI

PROMPTS = [
    "Explain how a hash map handles collisions, with a short Python example.",
    "Write a 300-word product description for a waterproof hiking backpack.",
    "Summarise the causes of the 2008 financial crisis in five bullet points, each with one sentence of detail.",
    "Write a Python function that parses ISO-8601 durations like P1DT2H30M into seconds, with tests.",
    "Compare TCP and QUIC for a mobile app developer: latency, reliability, and deployment.",
    "Draft a polite email to a landlord asking to fix a broken heater, then a shorter follow-up.",
]
MAX_TOKENS = 512
# Qwen's recommended sampling for non-thinking answers; the same on every target.
SAMPLING = {"temperature": 0.7, "top_p": 0.8}


def one(client, model, prompt, extra):
    start = time.perf_counter()
    first = None
    tokens = 0
    stream = client.chat.completions.create(
        model=model, messages=[{"role": "user", "content": prompt}], max_tokens=MAX_TOKENS,
        stream=True, stream_options={"include_usage": True}, extra_body=extra, **SAMPLING)
    for chunk in stream:
        delta = chunk.choices[0].delta if chunk.choices else None
        if delta and (delta.content or getattr(delta, "reasoning", None)) and first is None:
            first = time.perf_counter()
        if chunk.usage:
            tokens = chunk.usage.completion_tokens
    end = time.perf_counter()
    if first is None or tokens < 2:
        raise RuntimeError("no tokens streamed")
    return {"ttft_ms": (first - start) * 1000, "tok_s": (tokens - 1) / (end - first), "tokens": tokens}


def race(name, client, model, extra, rounds):
    runs, errors = [], 0
    for _ in range(rounds):
        for p in PROMPTS:
            try:
                runs.append(one(client, model, p, extra))
            except Exception as e:  # a failed request is counted, not hidden
                errors += 1
                print(f"  {name}: {type(e).__name__}: {str(e)[:120]}")
    if not runs:
        return {"target": name, "runs": 0, "errors": errors}
    tps = sorted(r["tok_s"] for r in runs)
    ttft = sorted(r["ttft_ms"] for r in runs)
    p90 = lambda xs, hi: xs[int(len(xs) * 0.9)] if hi else xs[int(len(xs) * 0.1)]
    return {"target": name, "runs": len(runs), "errors": errors,
            "tok_s_p50": statistics.median(tps), "tok_s_p10": p90(tps, False),
            "ttft_p50": statistics.median(ttft), "ttft_p90": p90(ttft, True),
            "tokens_mean": statistics.mean(r["tokens"] for r in runs)}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--denv-model", help="model name on the denv endpoint; omit to skip denv")
    ap.add_argument("--providers", default="", help="OpenRouter provider slugs, comma-separated")
    ap.add_argument("--or-model", default="qwen/qwen3.8-27b")
    ap.add_argument("--rounds", type=int, default=2)
    ap.add_argument("--out", default="race.json")
    a = ap.parse_args()

    results = []
    if a.denv_model:
        c = OpenAI(base_url=os.environ["DENV_BASE_URL"], api_key=os.environ["DENV_API_KEY"])
        results.append(race("denv", c, a.denv_model,
                            {"chat_template_kwargs": {"enable_thinking": False}}, a.rounds))
    slugs = [s for s in a.providers.split(",") if s]
    if slugs:
        c = OpenAI(base_url="https://openrouter.ai/api/v1", api_key=os.environ["OPENROUTER_API_KEY"])
        for slug in slugs:
            results.append(race(slug, c, a.or_model,
                                {"provider": {"order": [slug], "allow_fallbacks": False},
                                 "reasoning": {"enabled": False}}, a.rounds))

    print(f"\n{'target':<14}{'runs':>5}{'err':>5}{'tok/s p50':>11}{'tok/s p10':>11}{'TTFT p50':>10}{'TTFT p90':>10}{'avg toks':>10}")
    for r in sorted(results, key=lambda r: -r.get("tok_s_p50", 0)):
        if not r["runs"]:
            print(f"{r['target']:<14}{0:>5}{r['errors']:>5}")
            continue
        print(f"{r['target']:<14}{r['runs']:>5}{r['errors']:>5}{r['tok_s_p50']:>11.1f}{r['tok_s_p10']:>11.1f}"
              f"{r['ttft_p50']:>8.0f}ms{r['ttft_p90']:>8.0f}ms{r['tokens_mean']:>10.0f}")
    with open(a.out, "w") as f:
        json.dump({"at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), "max_tokens": MAX_TOKENS,
                   "sampling": SAMPLING, "results": results}, f, indent=1)


if __name__ == "__main__":
    main()