#!/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, so
what matters is whether that prefix is served from the cache or re-prefilled.
This script drives that shape and prints TTFT per turn.

Usage:
    bench-turn-latency.py <base-url> <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 the prompt itself and
therefore llama.cpp's prefix-cache key, which is what keeps a run from silently
inheriting a warm cache from the previous one.

One endpoint: send against it directly (:7770). The routing measurement this
harness originally produced (round-robin 2676 ms vs pinned 1096 ms warm turns on
90k-token contexts, 2026-09-16) belonged to the retired two-replica proxy and is
kept in the runbook as history. The current deployment is ONE process with ONE
slot, so concurrent families queue rather than spread — the harness reports that
shape honestly, and a saturated slot shows up as first-turn TTFT that includes
the wait.
"""

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())