#!/usr/bin/env python3
"""Benchmark all llama-swap models: prompt eval + generation speed."""
import json
import time
import urllib.request
import urllib.error
import sys
import os

BASE_URL = "http://127.0.0.1:9292/v1"

MODELS = [
    ("heretic-35b",      "heretic-35b Q4_K_M (daily driver, MoE 35B MTP)"),
    ("qwopus-9b",        "qwopus-9b Q6_K (coder, MoE 9B MTP)"),
    ("ds-flash-9b",      "ds-flash-9b Q6_K (Flash, MoE 9B MTP)"),
    ("qwen35-9b-ud",     "qwen35-9b-ud Q5_K_XL (uncensored, dense 9B)"),
    ("gemma-26b",        "gemma-26b Q4_K_XL (multimodal 26B)"),
    ("qwen36-35b-ud",    "qwen36-35b-ud Q4_K_M (backup UD, MoE 35B)"),
]

LONG_PROMPT = """Explain in detail the architecture, training procedure, and mathematical foundations of transformer neural networks. Cover the following aspects comprehensively:
1. The attention mechanism: scaled dot-product attention, multi-head attention, why scaling by sqrt(d_k) is necessary
2. Positional encoding: sinusoidal vs learned, the mathematical formula for sinusoidal PE, why it allows the model to handle sequence order
3. The encoder stack: layer normalization, residual connections, feed-forward networks (FFN), the role of each component
4. The decoder stack: masked self-attention, cross-attention, how it differs from encoder
5. Training details: teacher forcing, label smoothing, learning rate scheduling with warmup, the Adam optimizer configuration
6. Loss function: cross-entropy loss, why we predict the next token, how padding masks work
7. Inference: autoregressive decoding, beam search vs sampling, temperature scaling, top-k and top-p filtering
8. Scaling laws: how performance scales with model size, dataset size, and compute budget
9. Variants and improvements: RoPE, ALiBi, GQA, MQA, SwiGLU, RMSNorm, Flash Attention
10. Practical considerations: memory usage during training and inference, KV cache, continuous batching, tensor parallelism

For each point, include the specific equations, hyperparameter choices, and design rationale."""

SHORT_PROMPT = "Write a comprehensive 500-word essay on the future implications of artificial general intelligence for human society, covering economic, ethical, and social dimensions."


def api_call(model_id, messages, max_tokens, temperature=0.0, timeout=300):
    """Make a non-streaming API call and return parsed response + timing."""
    data = json.dumps({
        "model": model_id,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": temperature,
    }).encode()

    req = urllib.request.Request(
        f"{BASE_URL}/chat/completions",
        data=data,
        headers={"Content-Type": "application/json"},
        method="POST"
    )

    start = time.time()
    resp = urllib.request.urlopen(req, timeout=timeout)
    body = resp.read()
    elapsed = time.time() - start

    result = json.loads(body)
    usage = result.get("usage", {})
    return {
        "time": elapsed,
        "prompt_tokens": usage.get("prompt_tokens", 0),
        "completion_tokens": usage.get("completion_tokens", 0),
    }


def benchmark_model(model_id, label):
    """Benchmark a single model. Returns dict of results or error."""
    print(f"\n{'='*60}")
    print(f"  {label}")
    print(f"{'='*60}")

    result = {"model": model_id, "label": label}

    # Warm-up / load
    print(f"  Warming up...", flush=True)
    try:
        w = api_call(model_id, [{"role": "user", "content": "Reply with exactly: ready"}], max_tokens=5)
        print(f"  Warm-up: {w['time']:.1f}s, {w['prompt_tokens']}+{w['completion_tokens']} tok")
    except Exception as e:
        result["error"] = f"Warm-up failed: {e}"
        print(f"  ❌ {result['error']}")
        return result

    time.sleep(3)

    # Bench 1: Prompt eval speed (long prompt, tiny generation)
    print(f"  [Bench 1] Prompt eval: long prompt with max_tokens=5...", flush=True)
    try:
        b1 = api_call(model_id, [{"role": "user", "content": LONG_PROMPT}], max_tokens=5)
        pe_speed = b1["prompt_tokens"] / b1["time"]
        result["bench1_time"] = round(b1["time"], 2)
        result["bench1_prompt_tok"] = b1["prompt_tokens"]
        result["bench1_comp_tok"] = b1["completion_tokens"]
        result["prompt_eval_tok_s"] = round(pe_speed, 1)
        print(f"    Prompt: {b1['prompt_tokens']} tok, Gen: {b1['completion_tokens']} tok, Time: {b1['time']:.1f}s")
        print(f"    → Prompt eval: {pe_speed:.1f} tok/s")
    except Exception as e:
        result["error_pe"] = str(e)
        print(f"    ❌ {e}")

    time.sleep(3)

    # Bench 2: Generation speed (short prompt, long generation)
    print(f"  [Bench 2] Generation: short prompt with max_tokens=500...", flush=True)
    try:
        b2 = api_call(model_id, [{"role": "user", "content": SHORT_PROMPT}], max_tokens=500, temperature=0.7)
        gen_speed = b2["completion_tokens"] / b2["time"]
        result["bench2_time"] = round(b2["time"], 2)
        result["bench2_prompt_tok"] = b2["prompt_tokens"]
        result["bench2_comp_tok"] = b2["completion_tokens"]
        result["gen_tok_s"] = round(gen_speed, 1)
        print(f"    Prompt: {b2['prompt_tokens']} tok, Gen: {b2['completion_tokens']} tok, Time: {b2['time']:.1f}s")
        print(f"    → Generation: {gen_speed:.1f} tok/s")
    except Exception as e:
        result["error_gen"] = str(e)
        print(f"    ❌ {e}")

    print(f"  ✅ Done")
    return result


def main():
    print(f"LLAMA-SWAP MODEL BENCHMARK")
    print(f"Config: {time.strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"{'='*60}\n")

    results = []
    for model_id, label in MODELS:
        r = benchmark_model(model_id, label)
        results.append(r)
        with open("/tmp/benchmark_results.json", "w") as f:
            json.dump(results, f, indent=2)
        print()

    # Summary table
    print(f"\n{'='*60}")
    print(f"  SUMMARY")
    print(f"{'='*60}")
    print(f"{'Model':<35} {'Prompt Eval':<15} {'Generation':<15} {'Status'}")
    print(f"{'-'*75}")
    for r in results:
        pe = f"{r.get('prompt_eval_tok_s', '--')} t/s" if r.get('prompt_eval_tok_s') else "--"
        gen = f"{r.get('gen_tok_s', '--')} t/s" if r.get('gen_tok_s') else "--"
        status = "✅" if not r.get("error") else "❌"
        name = r.get("label", r["model"])[:35]
        print(f"{name:<35} {pe:<15} {gen:<15} {status}")

    print(f"\nResults saved to /tmp/benchmark_results.json")


if __name__ == "__main__":
    main()