#!/usr/bin/env python3 -u
"""Benchmark all fleet models through llama-swap. Single pass per model: cold load + gen + gen2 + PP + VRAM."""
import json, os, re, subprocess, sys, time, urllib.request, urllib.error

BASE = "http://127.0.0.1:9292/v1"
GEN_PROMPT = "Write a detailed Python script implementing a markdown parser with support for headers, lists, code blocks, and bold/italic text."
PP_PROMPT  = "What are the key differences between Mixture of Experts and Dense transformer architectures in large language models? Explain the routing mechanism, training stability concerns, and inference efficiency trade-offs."

def api(model, prompt, max_tok, temp=0.0):
    data = json.dumps({"model": model, "messages": [{"role": "user", "content": prompt}], "max_tokens": max_tok, "temperature": temp}).encode()
    req = urllib.request.Request(f"{BASE}/chat/completions", data=data, headers={"Content-Type": "application/json"})
    t0 = time.time()
    resp = urllib.request.urlopen(req, timeout=300)
    body = resp.read()
    elapsed = time.time() - t0
    r = json.loads(body)
    u = r.get("usage", {})
    return elapsed, u.get("prompt_tokens", 0), u.get("completion_tokens", 0)

def vram():
    r = subprocess.run(["nvidia-smi", "--query-gpu=memory.used,memory.free", "--format=csv,noheader,nounits"], capture_output=True, text=True, timeout=5)
    parts = r.stdout.strip().split(", ")
    return int(parts[0]), int(parts[1])

def benchmark_one(mid, label):
    print(f"\n{'─'*75}")
    print(f"  {label}")
    vram_before, _ = vram()
    t0 = time.time()
    t_warm, _, _ = api(mid, "say ok", 2)
    load_time = time.time() - t0
    vram_u, vram_f = vram()
    time.sleep(2)

    t_gen, _, ct = api(mid, GEN_PROMPT, 500, 0.7)
    gen1 = ct / t_gen if t_gen else 0
    time.sleep(1)

    t_gen2, _, ct2 = api(mid, GEN_PROMPT, 500, 0.7)
    gen2 = ct2 / t_gen2 if t_gen2 else 0
    time.sleep(1)

    t_pp, pt, _ = api(mid, PP_PROMPT, 5)
    pp = pt / t_pp if t_pp else 0

    vram_model = vram_u - vram_before
    print(f"  Load: {load_time:.0f}s | Gen: {gen1:.1f} t/s | Gen2: {gen2:.1f} t/s | PP: {pp:.0f} t/s")
    print(f"  VRAM: {vram_model:+d} MiB (used {vram_u} / free {vram_f})")

    return {
        "model": mid, "cold_load_s": round(load_time, 1),
        "gen_tok_s": round(gen1, 1), "gen2_tok_s": round(gen2, 1),
        "pp_tok_s": round(pp, 1), "vram_used": vram_u, "vram_free": vram_f,
    }

if __name__ == "__main__":
    models = [
        ("qwythos-9b-mtp-q6", "Qwythos-9B-MTP-Q6 (Q6_K, 131K)"),
        ("ornith-9b-q6",      "Ornith-9B-Q5 (Q5_K_M, 64K)"),
        ("qwen36-35b-mtp",    "Qwen3.6-35B-MTP (Q4_K_M, 230K)"),
        ("gemma-26b-200k",    "Gemma-26B-200K (Q4_K_XL, 230K)"),
        ("glm-4.7-flash",     "GLM-4.7-Flash (Q4_K_M, 200K)"),
        ("ornith-35b-q6-mtp", "Ornith-35B-Q6-MTP (Q6_K, 131K)"),
    ]
    print(f"{'='*75}\n  FLEET BENCHMARK — {time.strftime('%Y-%m-%d %H:%M:%S')}\n{'='*75}")
    results = [benchmark_one(mid, label) for mid, label in models]
    print(f"\n{'='*75}\n  SUMMARY\n{'='*75}")
    hdr = f"{'Model':<22} {'Load':>5} {'Gen':>7} {'Gen2':>7} {'PP':>8} {'VRAM':>6} {'Free':>6}"
    print(hdr); print('-'*len(hdr))
    for r in results:
        print(f"{r['model']:<22} {r['cold_load_s']:>5.0f}s {r['gen_tok_s']:>7.1f} {r['gen2_tok_s']:>7.1f} {r['pp_tok_s']:>8.0f} {r['vram_used']:>6} {r['vram_free']:>6}")
    with open("/tmp/fleet-benchmark.json", "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nSaved to /tmp/fleet-benchmark.json")