#!/usr/bin/env python3
"""bench-conversation.py — what an endpoint actually feels like, in tokens/s.

The other harnesses here answer engine questions (which batch, which proposer,
which cache). This one answers the user question: give the same ten jobs to an
OpenAI-compatible endpoint and report the rate it answers them at. Every job is
a real conversation turn — system prompt plus one user turn — because that is
the shape the fleet serves, and it is measured by STREAMING so prefill and
decode are separated:

  ttft          seconds to the first content delta  -> prefill + queue
  decode tok/s  completion_tokens / (total - ttft)  -> the rate you read at

Both arms of a comparison must run the SAME prompt bytes, so a per-run salt
rides in the system prompt: it changes the prompt itself, which is llama.cpp's
prefix-cache key and vLLM's APC key, so no arm silently inherits a warm cache
from the one before it. Run one arm, restart the engine, run the next with a
new salt.

Answers are generated at temperature 0 so the arms can be compared for
identity as well as speed. Identity is only meaningful across configurations
of the SAME model — different runtimes serving different quants legitimately
answer differently, which is why the first line of every answer is printed.

Usage:
    bench-conversation.py <base-url> [--model NAME] [--salt TAG]
                          [--max-tokens N] [--burst N] [--out file.json]

    bench-conversation.py http://127.0.0.1:7770 --model frank-medium --salt L1
    bench-conversation.py http://127.0.0.1:7770 --model frank-medium --salt R2 --burst 4

One engine at a time: on a single GPU pair there is no room for two copies of a
27B model, and the benchmarking policy requires the production service to be
stopped while anything measures.
"""

import argparse
import json
import math
import statistics
import sys
import threading
import time
import urllib.error
import urllib.request

# Ten jobs that cover what this fleet is actually asked to do. Ordered roughly
# by prefill size so a runtime that is good at short prompts and one that is
# good at long ones cannot hide behind an average.
SYSTEM = (
    "You are a staff engineer on a homelab fleet. Be concise and concrete. "
    "Salt: {salt}"
)

BATTERY = [
    ("greet", 128, "Say hello in one sentence."),
    ("shell", 192, "One bash command: find every file under /home modified in "
     "the last 2 days, sorted by size, largest first."),
    ("math", 256, "A train covers 180 km in 1h40m, then 240 km in 2h20m. "
     "What is the average speed in km/h? Show the arithmetic."),
    ("json", 384, "Emit exactly 12 JSON records, one per line, no prose, each "
     'with keys "host", "service", "port" and "up" for hosts dankclaw, '
     "dankserver and dankware running llama-server, caddy and grafana."),
    ("review", 512, "Review this function for bugs and name the failure mode:\n"
     "def rotate(cfg, n):\n    for k in cfg:\n        cfg[k] = cfg[k][n:] + cfg[k][:n]\n"
     "    return cfg\n"),
    ("explain", 448, "Explain in four sentences why PCIe atomic operation routing "
     "matters for GPU collectives, and what fails when a bridge reports "
     "AtomicOpsCap: Routing-."),
    ("translate", 256, "Translate to German, then to Japanese, one line each: "
     "The load balancer refused the connection because the port was still held."),
    ("plan", 640, "Give a 6-step runbook for swapping the model behind a live "
     "LLM endpoint on a host that serves real users, with a rollback at each step."),
    ("debug", 512, "A systemd user service restarts every 10 seconds and its log "
     "ends with 'Error starting server failed to bind udp listener on 10.89.1.1:53: "
     "Cannot assign requested address'. What is the root cause and the repair?"),
    ("longdoc", 512, None),  # filled from LONG_FILES at runtime
]

# The long-context job reads real repo files so prefill is large and identical
# across arms (~30k characters, roughly 8k tokens).
LONG_FILES = [
    "/home/dank/serverAppConfigs/AGENTS.md",
    "/home/dank/serverAppConfigs/serverAppConfigs-docs/LLM/"
    "frank-medium-llamacpp-dflash2.md",
]
LONG_CHARS = 30_000


def long_prompt() -> str:
    parts = []
    for path in LONG_FILES:
        try:
            with open(path, encoding="utf-8", errors="replace") as handle:
                parts.append(handle.read())
        except OSError:
            continue
    text = "\n\n".join(parts)[:LONG_CHARS]
    return ("Using only this document, answer: which deployment shapes for "
            "two GPUs are ruled out by missing PCIe atomics, and what is the "
            "one software path that still lets both cards serve one model? "
            "Cite the measured numbers.\n\n---\n" + text)


def stream_chat(base: str, model: str, messages, max_tokens: int, timeout: int):
    """One streamed chat completion -> (ttft, total, completion_tokens, text).

    Returns ttft=None when the engine produced no content delta.
    """
    body = {
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": 0,
        "stream": True,
        "stream_options": {"include_usage": True},
    }
    request = urllib.request.Request(
        base + "/v1/chat/completions",
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"},
    )
    start = time.perf_counter()
    ttft = None
    deltas = 0
    usage_tokens = None
    text = []
    with urllib.request.urlopen(request, timeout=timeout) as response:
        for raw in response:
            line = raw.decode("utf-8", "replace").strip()
            if not line.startswith("data:"):
                continue
            payload = line[5:].strip()
            if payload == "[DONE]":
                break
            try:
                chunk = json.loads(payload)
            except json.JSONDecodeError:
                continue
            if chunk.get("usage"):
                usage_tokens = chunk["usage"].get("completion_tokens")
            for choice in chunk.get("choices", []):
                piece = (choice.get("delta") or {})
                content = piece.get("content") or piece.get("reasoning_content") or ""
                if content:
                    if ttft is None:
                        ttft = time.perf_counter() - start
                    deltas += 1
                    text.append(content)
    total = time.perf_counter() - start
    return ttft, total, (usage_tokens or deltas), "".join(text)


def usage_prompt(base: str, model: str, messages, max_tokens: int):
    """prompt_tokens for a job, from one non-streamed call with max_tokens=1."""
    body = {
        "model": model,
        "messages": messages,
        "max_tokens": 1,
        "temperature": 0,
        "stream": False,
    }
    request = urllib.request.Request(
        base + "/v1/chat/completions",
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(request, timeout=300) as response:
        return json.load(response).get("usage", {}).get("prompt_tokens")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("base")
    ap.add_argument("--model", default="frank-medium")
    ap.add_argument("--salt", default=f"C{int(time.time())}")
    ap.add_argument("--max-tokens", type=int, default=None)
    ap.add_argument("--burst", type=int, default=0,
                    help="also run N concurrent copies of the plan job")
    ap.add_argument("--timeout", type=int, default=1800)
    ap.add_argument("--out", default=None)
    args = ap.parse_args()

    base = args.base.rstrip("/")
    rows = []
    print(f"# battery salt={args.salt} model={args.model} base={base}\n")
    print(f"{'job':<10} {'ptok':>7} {'ctok':>6} {'ttft_s':>8} {'dec_tps':>8} "
          f"{'total_s':>8}  answer")

    for name, cap, prompt in BATTERY:
        if prompt is None:
            prompt = long_prompt()
        messages = [
            {"role": "system", "content": SYSTEM.format(salt=args.salt)},
            {"role": "user", "content": prompt},
        ]
        limit = args.max_tokens or cap
        try:
            ptok = usage_prompt(base, args.model, messages, limit)
            ttft, total, ctok, text = stream_chat(
                base, args.model, messages, limit, args.timeout)
        except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError) as exc:
            print(f"{name:<10} FAILED {exc}")
            rows.append({"job": name, "error": str(exc)})
            continue
        decode = None
        if ttft is not None and ctok and total > ttft:
            decode = (ctok - 1) / (total - ttft)
        prefill_rate = (ptok / ttft) if (ptok and ttft) else None
        first = text.strip().splitlines()[0][:48] if text.strip() else "(empty)"
        print(f"{name:<10} {ptok or 0:>7} {ctok:>6} "
              f"{(ttft or 0):>8.2f} {(decode or 0):>8.1f} {total:>8.2f}  {first}")
        rows.append({"job": name, "prompt_tokens": ptok, "completion_tokens": ctok,
                     "ttft_s": ttft, "decode_tps": decode, "prefill_tps": prefill_rate,
                     "total_s": total, "answer_head": first})

    good = [r for r in rows if r.get("decode_tps")]
    if good:
        rates = [r["decode_tps"] for r in good]
        ttfts = [r["ttft_s"] for r in good]
        pre = [r["prefill_tps"] for r in good if r.get("prefill_tps")]
        print(f"\ngeomean decode  {statistics.geometric_mean(rates):8.1f} tok/s"
              f"   (min {min(rates):.1f} / max {max(rates):.1f})")
        print(f"mean ttft       {statistics.mean(ttfts):8.2f} s")
        if pre:
            print(f"geomean prefill {statistics.geometric_mean(pre):8.1f} tok/s")

    if args.burst > 1:
        prompt = next(p for n, _, p in BATTERY if n == "plan")
        messages = [
            {"role": "system", "content": SYSTEM.format(salt=args.salt + "B")},
            {"role": "user", "content": prompt},
        ]
        results = []
        lock = threading.Lock()

        def worker():
            try:
                ttft, total, ctok, _ = stream_chat(
                    base, args.model, messages, 640, args.timeout)
                with lock:
                    results.append((ttft, total, ctok))
            except Exception as exc:  # noqa: BLE001 - report, do not abort the burst
                with lock:
                    results.append((None, 0.0, str(exc)))

        threads = [threading.Thread(target=worker) for _ in range(args.burst)]
        t0 = time.perf_counter()
        for thread in threads:
            thread.start()
        for thread in threads:
            thread.join()
        wall = time.perf_counter() - t0
        ok = [r for r in results if r[0] is not None]
        if ok:
            tokens = sum(r[2] for r in ok)
            print(f"burst {args.burst}x plan: wall {wall:.1f}s  aggregate "
                  f"{tokens / wall:.1f} tok/s  "
                  f"per-stream {statistics.mean([r[2] / r[1] for r in ok]):.1f} tok/s"
                  f"  mean ttft {statistics.mean([r[0] for r in ok]):.2f}s")
        else:
            print(f"burst {args.burst}x plan: all failed: {results[0][2]}")

    if args.out:
        with open(args.out, "w", encoding="utf-8") as handle:
            json.dump({"salt": args.salt, "model": args.model, "base": base,
                       "rows": rows, "burst": args.burst}, handle, indent=1)
        print(f"wrote {args.out}")
    return 0 if good else 1


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