#!/usr/bin/env python3
"""
Decode throughput benchmark for frank-default-large.
Starts the server with a given --threads-batch value, sends an 8K+ token prompt,
generates 256 tokens, measures generation tok/s.

Each --threads-batch value requires a cold server load (~3 min).
"""
import json, time, urllib.request, urllib.error, subprocess, os, sys

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"
PORT = 7770
N_GEN = 256
PROMPT_FILE = "/tmp/bench-prompts/prompt_code.txt"
ENV_FILE = os.path.join(os.path.dirname(__file__), "frank-default-large.env")
SERVER_DIR = os.path.dirname(SERVE)
ROCM_PATH = "/opt/rocm"

def make_env():
    """Build the same env as serve-frank-large.sh."""
    env = {**os.environ, "HSA_OVERRIDE_GFX_VERSION": "11.5.1", "HIP_VISIBLE_DEVICES": "0"}
    
    # Source env file for tunables
    if os.path.exists(ENV_FILE):
        with open(ENV_FILE) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    k, v = line.split("=", 1)
                    env[k] = v
    
    # Build LD_LIBRARY_PATH
    lib_path = f"{SERVER_DIR}:{ROCM_PATH}/lib:{ROCM_PATH}/lib/rocm_sysdeps/lib:{ROCM_PATH}/lib/llvm/lib"
    if os.environ.get("LD_LIBRARY_PATH"):
        lib_path += f":{os.environ['LD_LIBRARY_PATH']}"
    env["LD_LIBRARY_PATH"] = lib_path
    env["ROCM_PATH"] = ROCM_PATH
    
    return env

ENV = make_env()

def kill_all():
    subprocess.run(["pkill", "-9", "llama-server"], 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)
            if json.loads(r.read()).get("status") == "ok":
                return True
        except:
            pass
        time.sleep(1)
    return False

def send_request(prompt, max_tokens=N_GEN):
    """Send chat completion, measure generation throughput."""
    body = json.dumps({
        "model": "frank-default-large",
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "stream": False,
        "temperature": 0.7,
        "top_p": 0.8
    }).encode()
    
    t0 = time.time()
    req = urllib.request.Request(
        f"http://127.0.0.1:{PORT}/v1/chat/completions",
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST"
    )
    r = urllib.request.urlopen(req, timeout=1200)
    resp = json.loads(r.read())
    wall = time.time() - t0
    
    usage = resp.get("usage", {})
    pt = usage.get("prompt_tokens", 0)
    gt = usage.get("completion_tokens", 0)
    gtps = gt / wall if wall > 0 else 0
    return pt, gt, wall, gtps

def load_and_extend_prompt():
    """Load prompt and extend to ~8K tokens if needed."""
    with open(PROMPT_FILE) as f:
        text = f.read()
    
    # Count approximate tokens (rough: ~4 chars per token for code)
    approx_tokens = len(text) / 4.0
    needed = 8200
    
    if approx_tokens < needed:
        repeats = needed // max(int(approx_tokens), 1) + 1
        text = (text + "\n") * repeats
        words = text.split()[:needed]
        text = " ".join(words)
    
    return text

def main():
    tb_values = [int(x) for x in sys.argv[1:]] if len(sys.argv) > 1 else [16, 32, 64, 128]
    results = []
    
    prompt = load_and_extend_prompt()
    n_words = len(prompt.split())
    print(f"Prompt: {n_words} words (~{n_words} tokens), {len(prompt)} chars\n", flush=True)
    
    for tb in tb_values:
        print(f"\n{'='*60}")
        print(f"threads-batch = {tb}")
        print(f"{'='*60}", flush=True)
        
        kill_all()
        print("  Starting server...", flush=True)
        
        lib_path = ENV.get("LD_LIBRARY_PATH", "")
        cmd = [
            "/usr/bin/env",
            "HSA_OVERRIDE_GFX_VERSION=11.5.1",
            "HIP_VISIBLE_DEVICES=0",
            f"LD_LIBRARY_PATH={lib_path}",
            f"ROCM_PATH={ROCM_PATH}",
            SERVE, "--model", MODEL, "--port", str(PORT), "--jinja",
            "--device", "ROCm0", "--gpu-layers", "999", "--fit", "off",
            "--threads", "16", "--threads-batch", str(tb),
            "--parallel", "2", "--ctx-size", "131072",
            "--batch-size", "4096", "--ubatch-size", "4096",
            "--flash-attn", "on",
            "--cache-type-k", "f16", "--cache-type-v", "f16",
            "--cache-prompt", "--cache-idle-slots", "--metrics"
        ]
        
        srv = subprocess.Popen(cmd, env=ENV, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        
        ready = wait_ready()
        if not ready:
            print("  FAILED: server did not become ready", flush=True)
            srv.terminate()
            srv.wait()
            for line in (srv.stdout.readlines() or [])[-20:]:
                print(f"  LOG: {line.rstrip()}", flush=True)
            continue
        
        print("  Server ready.", flush=True)
        
        # Warmup (small request)
        try:
            send_request("Hi", 16)
        except Exception as e:
            print(f"  Warmup failed: {e}", flush=True)
            srv.terminate()
            srv.wait()
            continue
        
        # Benchmark
        print(f"  Generating {N_GEN} tokens...", flush=True)
        try:
            pt, gt, wall, gtps = send_request(prompt, N_GEN)
            print(f"  {gt}/{N_GEN} gen tokens in {wall:.1f}s = {gtps:.1f} tok/s")
            print(f"  Prompt tokens: {pt}")
            results.append({"tb": tb, "wall": wall, "gtps": gtps, "gt": gt, "pt": pt})
        except Exception as e:
            print(f"  FAILED: {e}", flush=True)
        
        srv.terminate()
        srv.wait()
        time.sleep(2)
    
    # Summary
    print(f"\n{'='*60}")
    print("SUMMARY")
    print(f"{'='*60}")
    for r in results:
        print(f"  tb={r['tb']:>3}: {r['gt']} gen tokens, {r['wall']:.1f}s, {r['gtps']:.1f} tok/s")
    
    if results:
        best = max(results, key=lambda r: r['gtps'])
        print(f"\nBest: tb={best['tb']} at {best['gtps']:.1f} tok/s")
    
    with open("/tmp/bench_decode_results.txt", "w") as f:
        f.write("tb,gen_tokens,wall_s,gen_tps,prompt_tokens\n")
        for r in results:
            f.write(f"{r['tb']},{r['gt']},{r['wall']:.1f},{r['gtps']:.1f},{r['pt']}\n")
    
    print(f"\nResults: /tmp/bench_decode_results.txt")

if __name__ == "__main__":
    main()