import subprocess, time, os, sys, csv, io

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"
BENCH = "/home/dank/GitPackages/ROCmFPX/build-qwen38/bin/llama-bench"
TB_LIST = [16, 32, 48, 64, 80, 96, 128]
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():
    import urllib.request
    for i in range(600):
        try:
            r = urllib.request.urlopen("http://127.0.0.1:7770/health", timeout=2)
            if r.status == 200:
                return True
        except:
            pass
        time.sleep(1)
    return False

def parse_bench_output(stdout, stderr):
    """Parse llama-bench CSV output. Columns (0-indexed): n_prompt=33, avg_ts=39"""
    prefill = decode = None
    ctx = "?"
    try:
        reader = csv.reader(io.StringIO(stdout))
        header = next(reader, None)
        if header is None:
            return None, None, "?"
        
        # Find prefill row: n_prompt matches our prompt size, n_gen matches our gen size
        # For n-gen=0, we get prefill only. The prefill row has n_prompt > 0.
        for row in reader:
            if len(row) >= 41:
                try:
                    n_prompt = int(row[33])
                    n_gen = int(row[34])
                    avg_ts = float(row[39])  # avg_tokens_per_second
                    
                    if n_prompt > 0 and n_gen == 0:
                        # This is the prefill row
                        prefill = avg_ts
                        ctx = row[40].strip() if row[40].strip() else "?"
                        return prefill, None, ctx
                    elif n_prompt > 0 and n_gen > 0:
                        # Mixed prefill+decode: avg_ts is overall, use prefill row if available
                        if prefill is None:
                            prefill = avg_ts  # fallback
                except (ValueError, IndexError):
                    continue
    except Exception as e:
        print(f"  CSV parse error: {e}", flush=True)
    
    return prefill, decode, ctx

results = []
for tb in TB_LIST:
    print(f"=== threads-batch={tb} ===", flush=True)
    kill_all()

    srv = subprocess.Popen(
        [SERVE, "--model", MODEL, "--port", "7770", "--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()
        out = srv.stdout.read(1000)
        print(f"Server failed: {out[:300]}")
        sys.exit(1)
    print("Server ready, running bench...", flush=True)

    t0 = time.time()
    try:
        # Use n-gen=0 to get prefill-only throughput
        res = subprocess.run(
            [BENCH, "--model", MODEL, "--threads", "16",
             "--n-prompt", "65536", "--n-gen", "0",
             "--repetitions", "3", "--output", "csv",
             "--flash-attn", "on", "--n-gpu-layers", "999"],
            env=ENV, capture_output=True, text=True, timeout=600)
        stdout = res.stdout + res.stderr
    except subprocess.TimeoutExpired:
        print("Bench timed out!")
        srv.terminate()
        results.append({"tb": tb, "prefill": None, "decode": None, "ctx": None, "time": 600})
        continue

    elapsed = round(time.time() - t0, 1)
    print(f"Bench rc={res.returncode} stdout_len={len(res.stdout)} stderr_len={len(res.stderr)}", flush=True)
    if res.returncode != 0:
        print(f"STDERR: {res.stderr[:500]}", flush=True)
    if res.stdout.strip():
        print(f"CSV OUTPUT:\n{res.stdout[:300]}", flush=True)

    prefill, decode, ctx = parse_bench_output(res.stdout, res.stderr)

    row = {"tb": tb, "prefill": prefill, "decode": decode, "ctx": ctx, "time": elapsed}
    results.append(row)

    pt = f"{prefill:.1f}" if prefill else "N/A"
    dt = f"{decode:.1f}" if decode else "N/A"
    print(f"  prefill={pt} tok/s  decode={dt} tok/s  ctx={ctx}  time={elapsed}s", flush=True)

    srv.terminate()
    srv.wait()
    time.sleep(2)

print("\n==========================================", flush=True)
print("SUMMARY", flush=True)
print("==========================================", flush=True)
print(f"{'tb':>5}  {'prefill':>12}  {'decode':>12}  {'ctx':>10}  {'time':>8}", flush=True)
print("-" * 60, flush=True)
for r in results:
    pt = f"{r['prefill']:.1f}" if r["prefill"] else "N/A"
    dt = f"{r['decode']:.1f}" if r["decode"] else "N/A"
    print(f"{r['tb']:>5}  {pt:>12}  {dt:>12}  {str(r['ctx']):>10}  {r['time']:>8.0f}", flush=True)

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

with open("/tmp/bench_tb_results.txt", "w") as f:
    for r in results:
        pt = str(r["prefill"]) if r["prefill"] else "N/A"
        dt = str(r["decode"]) if r["decode"] else "N/A"
        f.write(f"{r['tb']},{pt},{dt},{r['ctx']},{r['time']}\n")