#!/usr/bin/env python3
"""
HTTP-based prefill-throughput sweep for frank-default-large.
Starts the server with a given --threads-batch value, sends a 64K-prompt via
the chat API (max_tokens=0), measures prompt_tokens/wall_time.

One request per value. Each --threads-batch requires a fresh cold server load
(~3 min), so repeating would add 2x overhead. The server also crashes on 2nd
request with max_tokens=0, making repeats impractical.
"""
import subprocess, time, os, sys, json, urllib.request, urllib.error

MODEL = "/home/dank/models/kingjones-Qwen3.8-Flash-Next-ROCmFP4-STRIX/Qwen3.8-Flash-Next-Q4_0-ROCmFP4-STRIX.gguf"
SERVE = "/home/dank/GitPackages/ROCmFPX/build-qwen38/bin/llama-server"
TB_LIST = [16, 32, 48, 64, 80, 96, 128]
PROMPT_FILE = "/tmp/bench-prompts/prompt_code.txt"
PORT = 7770
ENV = {**os.environ, "HSA_OVERRIDE_GFX_VERSION": "11.5.1", "HIP_VISIBLE_DEVICES": "0"}

def kill_all():
    subprocess.run(["pkill", "-9", "llama-server"], capture_output=True)
    subprocess.run(["pkill", "-9", "llama-bench"], capture_output=True)
    time.sleep(3)

def wait_ready():
    for i in range(600):
        try:
            r = urllib.request.urlopen(f"http://127.0.0.1:{PORT}/health", timeout=2)
            data = json.loads(r.read())
            if data.get("status") == "ok":
                return True
        except:
            pass
        time.sleep(1)
    return False

def send_prefill_request():
    """Send chat completion with max_tokens=0 (prefill only)."""
    with open(PROMPT_FILE, "r") as f:
        prompt = f.read()
    body = json.dumps({
        "model": "frank-default-large",
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": 1,
        "stream": False
    }).encode()
    req = urllib.request.Request(
        f"http://127.0.0.1:{PORT}/v1/chat/completions",
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST"
    )
    t0 = time.time()
    r = urllib.request.urlopen(req, timeout=600)
    resp = json.loads(r.read())
    wall = time.time() - t0
    usage = resp.get("usage", {})
    prompt_tokens = usage.get("prompt_tokens", 0)
    return prompt_tokens, wall

results = []
for tb in TB_LIST:
    print(f"=== threads-batch={tb} ===", flush=True)
    kill_all()
    
    # Start server
    print("  Starting server...", flush=True)
    srv = subprocess.Popen(
        [SERVE, "--model", MODEL, "--port", str(PORT), "--jinja",
         "--device", "ROCm0", "--gpu-layers", "999", "--fit", "off",
         "--threads", "16", "--threads-batch", str(tb),
         "--parallel", "2", "--ctx-size", "524288",
         "--batch-size", "4096", "--ubatch-size", "4096",
         "--flash-attn", "on",
         "--cache-type-k", "f16", "--cache-type-v", "f16",
         "--cache-prompt", "--cache-idle-slots", "--metrics"],
        env=ENV, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    
    ready = wait_ready()
    if not ready:
        srv.terminate()
        srv.wait()
        print("  FAILED: server didn't start")
        sys.exit(1)
    print("  Ready.", flush=True)
    
    # Single prefill request
    try:
        prompt_tokens, wall = send_prefill_request()
        if prompt_tokens > 0:
            tps = prompt_tokens / wall
            print(f"  {prompt_tokens} tok in {wall:.1f}s = {tps:.1f} tok/s", flush=True)
            results.append({"tb": tb, "avg": tps, "std": 0, "samples": 1})
        else:
            print(f"  FAILED (0 prompt_tokens, wall={wall:.1f}s)", flush=True)
            results.append({"tb": tb, "avg": None, "std": 0, "samples": 0})
    except Exception as e:
        print(f"  ERROR: {e}", flush=True)
        results.append({"tb": tb, "avg": None, "std": 0, "samples": 0})
    
    srv.terminate()
    srv.wait()
    time.sleep(2)

# Summary
print("\n" + "=" * 60)
print("SUMMARY")
print("=" * 60)
for r in results:
    if r["avg"] is not None:
        print(f"  tb={r['tb']:>3}  prefill={r['avg']:.1f} tok/s ({r['samples']} sample)")
    else:
        print(f"  tb={r['tb']:>3}  FAILED")

best = max((r for r in results if r["avg"] is not None), key=lambda r: r["avg"])
if best:
    print(f"\nBest: threads-batch={best['tb']} ({best['avg']:.1f} tok/s)")

with open("/tmp/bench_tb_results.txt", "w") as f:
    for r in results:
        avg = f"{r['avg']:.1f}" if r["avg"] is not None else "N/A"
        f.write(f"{r['tb']},{avg},0,{r['samples']}\n")