#!/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 = [
    ("qwen3.6-35b",        "qwen3.6-35b Q4_K_S (daily driver, MoE 35B)"),
    ("qwen3.6-mtp",        "qwen3.6-mtp Q4_K_S (MTP variant, MoE 35B)"),
    ("qwen3.5-9b",         "qwen3.5-9b Q4_K_M (uncensored, dense 9B)"),
    ("gemma-4-26b",        "gemma-4-26b IQ4_XS (multimodal 26B)"),
    ("qwen3-coder-30b",     "qwen3-coder-30b Q4_K_M (coding, MoE 30B)"),
    ("qwen2.5-coder-14b",  "qwen2.5-coder-14b Q4_K_M (dense 14B coder, 65K)"),
    ("qwen3-30b-2507",     "qwen3-30b-2507 Q4_K_M (MoE 30B July)"),
    ("qwen3.5-opus14",     "qwen3.5-opus14 MXFP4 (Opus distilled, MoE 14B)"),
]

# ~800-token prompt for eval benchmarking
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. This should be thorough enough for someone with a machine learning background to understand the full picture."""

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),
        "response_preview": result["choices"][0]["message"]["content"][:100] if result.get("choices") else "",
    }


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: 128K context across the board (131K Gemma, 65K Coder-14B dense)")
    print(f"MTP flags applied to original fleet (n-cpu-moe 128, no -fa off, batch 2048/512)")
    print(f"Time: {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)
        # Save incremental results
        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':<30} {'Context':<10} {'Prompt Eval':<15} {'Generation':<15} {'Status'}")
    print(f"{'-'*85}")
    for r in results:
        ctx = "128K" if "error" not in r or not r.get("error") else "FAIL"
        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"])[:30]
        print(f"{name:<30} {ctx:<10} {pe:<15} {gen:<15} {status}")

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

    # Memory note on coder14
    for r in results:
        if r["model"] == "qwen2.5-coder-14b" and not r.get("error"):
            print(f"\n⚠️  WARNING: qwen2.5-coder-14b (dense 14B) at 128K may be VRAM-constrained.")
            print(f"   If unstable, consider reverting to 65K or enabling --no-kv-offload.")


if __name__ == "__main__":
    main()