#!/usr/bin/env python3
"""bench-turn-latency.py — per-turn TTFT for an agent-loop shape.

Every other harness in this directory measures a *batch* of independent
requests. That answers throughput, and it is the wrong instrument for the number
a user feels: an agent loop re-sends the same large prefix turn after turn, and
the deployed routing pins such a conversation to one replica so its prefix cache
stays warm. This script drives that shape and prints TTFT per turn.

It is also the measurement that produced the routing decision on 2026-09-16
(round-robin 2676 ms mean warm turn against pinned 1096 ms on 90k-token
contexts), and the one that shows a cold-vs-warm prefill directly: pump the same
prompt twice and the second is a cache lookup.

Usage:
    bench-turn-latency.py <base-url> <minutes|turns> [families] [salt]
    bench-turn-latency.py http://127.0.0.1:7770 5 3 SALT1     # 3 families, 5 turns each

Prompts carry the salt as their FIRST bytes: that changes both the router's
identity window and vLLM's prefix-cache key, which is what keeps a run from
silently inheriting a warm cache from the previous one.

Two cards behind a proxy: send against the proxy (:7770) to measure the fleet,
or a replica directly (:7771, :7772) to measure one card.
"""

from __future__ import annotations

import json
import statistics
import sys
import threading
import time
import urllib.request

BASE = sys.argv[1].rstrip("/") if len(sys.argv) > 1 else "http://127.0.0.1:7770"
TURNS = int(sys.argv[2]) if len(sys.argv) > 2 else 5
FAMILIES = int(sys.argv[3]) if len(sys.argv) > 3 else 3
SALT = sys.argv[4] if len(sys.argv) > 4 else f"BENCH{int(time.time())}"

DOC = "/home/dank/serverAppConfigs/scripts/deploy-remote.sh"
MAX_CHARS = 120_000

lock = threading.Lock()
rows: list[tuple[str, float, bool]] = []


def ttft_ms(prompt: str) -> tuple[float, bool]:
    body = {
        "model": "frank-medium",
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": 32,
        "stream": True,
    }
    req = urllib.request.Request(
        BASE + "/v1/chat/completions",
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"},
    )
    start = time.time()
    try:
        with urllib.request.urlopen(req, timeout=900) as resp:
            for line in resp:
                if line.startswith(b"data: ") and b'"content"' in line:
                    return (time.time() - start) * 1000.0, True
    except Exception as exc:
        print(f"    error: {type(exc).__name__}: {str(exc)[:70]}", flush=True)
    return (time.time() - start) * 1000.0, False


def family(index: int) -> None:
    with open(DOC, errors="replace") as fh:
        doc = (fh.read() * 3)[:MAX_CHARS]
    # History grows the way a real conversation does, so later turns are
    # continuations — the requests the router pins and the cache can serve.
    history: list[dict] = []
    for turn in range(TURNS):
        prompt = (f"{SALT}-{index} conversation.\n{doc}\n"
                  f"Turn {turn}: which part of the above needs root?")
        ms, ok = ttft_ms(prompt)
        with lock:
            rows.append((f"fam{index}", ms, ok))
        print(f"  fam{index} turn{turn + 1}: {ms:9.1f} ms {'ok' if ok else 'FAIL'}",
              flush=True)
        history += [{"role": "assistant", "content": "Acknowledged."},
                    {"role": "user", "content": f"turn {turn}"}]


def main() -> int:
    threads = [threading.Thread(target=family, args=(i,), daemon=True)
               for i in range(FAMILIES)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    with lock:
        snapshot = list(rows)
    ok = [r for r in snapshot if r[2]]
    print(f"\nTURNS salt={SALT} families={FAMILIES} requests={len(snapshot)} "
          f"ok={len(ok)} failed={len(snapshot) - len(ok)}")
    for fam in sorted({r[0] for r in snapshot}):
        times = [r[1] for r in ok if r[0] == fam]
        if not times:
            continue
        warm = [t for t in times[1:] if t < 5000]
        print(f"  {fam}: first={times[0]:.0f} ms "
              f"median={statistics.median(times):.0f} ms "
              f"warm_turns={len(warm)}/{max(len(times) - 1, 1)}")
    return 0 if len(ok) == len(snapshot) else 1


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