"""Quick benchmark for running llama-server instance"""
import urllib.request, json, time, sys

BASE = "http://127.0.0.1:8000"
RUNS = 5
TOKENS = 200

def bench(max_tokens=TOKENS):
    payload = json.dumps({
        "model": "m",
        "messages": [{"role": "user", "content": "Count from 1 to 100, each number on a new line."}],
        "max_tokens": max_tokens,
        "temperature": 0
    }).encode()
    req = urllib.request.Request(
        f"{BASE}/v1/chat/completions",
        data=payload,
        headers={"Content-Type": "application/json"}
    )
    t0 = time.time()
    resp = json.loads(urllib.request.urlopen(req, timeout=300).read())
    dt = time.time() - t0
    ct = resp.get("usage", {}).get("completion_tokens", 0)
    return ct / dt if dt > 0 else 0, ct, dt

print("Warmup...", flush=True)
try:
    bench(20)
except Exception as e:
    print(f"Warmup failed: {e}")
    sys.exit(1)
print("Warmup done\n", flush=True)

speeds = []
for i in range(RUNS):
    tps, ct, dt = bench()
    speeds.append(tps)
    print(f"  Run {i+1}: {tps:.2f} t/s (tokens={ct}, time={dt:.2f}s)", flush=True)

avg = sum(speeds) / len(speeds)
best = max(speeds)
mn = min(speeds)
print(f"\n{'='*50}")
print(f"  RESULT: AVG {avg:.2f} / BEST {best:.2f} / MIN {mn:.2f} t/s")
print(f"{'='*50}")