#!/usr/bin/env python3
"""Benchmark a single llama.cpp model across multiple flag configurations.

Spawns llama-server directly (no llama-swap), runs warmup + 3 benchmarks
(prompt eval, generation, sustained generation), and saves results to JSON.

Designed for flag tuning: each config varies ONE parameter from baseline.
Kills and restarts the server between configs.

Usage:
  python3 benchmark_single_model.py

Edit the MODEL, BINARY, BASE_CMD, and CONFIGS globals at the top of this file
to match the model and configurations you want to test.
"""

import json
import time
import urllib.request
import urllib.error
import subprocess
import signal
import sys
import os
import re

# ============================================================================
# CONFIGURATION — Edit these for your model
# ============================================================================

MODEL = "/models/downloads/nex-agi_Nex-N2-mini-Q4_K_M.gguf"
BINARY = "/usr/local/bin/llama-server-upstream"
PORT = 5804  # must not conflict with any running llama-swap backend

BASE_CMD = [
    BINARY,
    "-m", MODEL,
    "--port", str(PORT),
    "--host", "127.0.0.1",
    "--jinja",
    '--chat-template-kwargs', '{"enable_thinking":false}',
    "-ngl", "99",
    "--no-mmap",
    "--mlock",
    "--prio", "2",
    "--poll", "80",
    "--timeout", "300",
    "--no-host",
    "-fitt", "512",
    "--metrics",
]

# Each config: name, desc, and extra args (appended to BASE_CMD).
# For new models, start with baseline then vary ONE variable at a time.
CONFIGS = [
    # Example configs — replace with yours
    {
        "name": "baseline",
        "desc": "250K ctx, q8_0/q8_0, 1024/256, t16/32, moe128, fa-on, cont-batch",
        "args": [
            "-c", "250000", "--ctx-size", "250000",
            "--cache-type-k", "q8_0", "--cache-type-v", "q8_0",
            "--batch-size", "1024", "--ubatch-size", "256",
            "--threads", "16", "--threads-batch", "32",
            "--n-cpu-moe", "128",
            "--flash-attn", "on",
            "--cont-batching",
            "--no-kv-offload",
            "--temp", "0.7", "--top-p", "0.95", "--top-k", "40",
        ],
    },
]

# ============================================================================
# Test prompts
# ============================================================================

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 essay on the future implications of artificial general intelligence for human society, covering economic, ethical, and social dimensions."

# ============================================================================
# Benchmark helpers
# ============================================================================

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

    req = urllib.request.Request(
        f"http://127.0.0.1:{PORT}/v1/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", {})
    timings = result.get("timings", {})

    return {
        "time": round(elapsed, 2),
        "prompt_tokens": usage.get("prompt_tokens", 0),
        "completion_tokens": usage.get("completion_tokens", 0),
        "prompt_ms": timings.get("prompt_ms", 0),
        "prompt_per_token_ms": timings.get("prompt_per_token_ms", 0),
        "predicted_ms": timings.get("predicted_ms", 0),
        "predicted_per_token_ms": timings.get("predicted_per_token_ms", 0),
        "prompt_per_second": timings.get("prompt_per_second", 0),
        "predicted_per_second": timings.get("predicted_per_second", 0),
    }


def wait_for_server(timeout=120):
    """Block until /health responds or timeout."""
    start = time.time()
    while time.time() - start < timeout:
        try:
            req = urllib.request.Request(f"http://127.0.0.1:{PORT}/health")
            urllib.request.urlopen(req, timeout=5)
            return True
        except Exception:
            time.sleep(1)
    return False


def kill_process(proc):
    """Gracefully kill a llama-server process."""
    if proc and proc.poll() is None:
        proc.send_signal(signal.SIGINT)
        try:
            proc.wait(timeout=10)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait()
        time.sleep(2)


def benchmark_config(config, results_file):
    """Run a single config: start server, warmup, 3 benchmarks, kill."""
    name = config["name"]
    desc = config["desc"]
    args = config["args"]

    print(f"\n{'='*70}")
    print(f"  [{name}] {desc}")
    print(f"{'='*70}")

    cmd = BASE_CMD + args
    proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    if not wait_for_server():
        print(f"  ❌ Server failed to start within timeout")
        kill_process(proc)
        return None

    print(f"  Server ready on port {PORT}")

    result = {"name": name, "desc": desc, "args": ' '.join(cmd[4:])}

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

    time.sleep(2)

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

    time.sleep(2)

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

    # Bench 3: Sustained generation
    time.sleep(2)
    print(f"  [Bench 3] Sustained gen: 200 tok...", flush=True)
    try:
        b3 = api_call([{"role": "user", "content": "Write a short story about a robot learning to paint."}], max_tokens=200, temperature=0.7)
        gen_speed3 = b3["completion_tokens"] / b3["time"]
        result["bench3"] = b3
        result["gen_sustained_tok_s"] = round(gen_speed3, 1)
        print(f"    Prompt: {b3['prompt_tokens']} tok, Gen: {b3['completion_tokens']} tok, Time: {b3['time']}s")
        print(f"    → Sustained gen: {gen_speed3:.1f} tok/s")
    except Exception as e:
        result["error_gen3"] = str(e)
        print(f"    ❌ {e}")

    # Cleanup
    kill_process(proc)

    # Save incremental results
    all_results = []
    if os.path.exists(results_file):
        with open(results_file) as f:
            all_results = json.load(f)
    all_results.append(result)
    with open(results_file, "w") as f:
        json.dump(all_results, f, indent=2)

    print(f"  ✅ Done")
    return result


def print_summary(results_file):
    """Print a formatted summary table from saved results."""
    with open(results_file) as f:
        results = json.load(f)

    print(f"\n{'='*100}")
    print(f"  BENCHMARK SUMMARY")
    print(f"{'='*100}")
    print(f"{'Config':<28} {'PE (t/s)':<12} {'PE (ms/t)':<12} {'Gen (t/s)':<12} {'Gen (ms/t)':<12} {'Sust. Gen':<12}")
    print(f"{'-'*28} {'-'*12} {'-'*12} {'-'*12} {'-'*12} {'-'*12}")

    for r in results:
        pe = f"{r.get('prompt_eval_tok_s', '--'):>8.1f}" if isinstance(r.get('prompt_eval_tok_s'), (int, float)) else f"{'--':>8}"
        pe_ms = f"{r.get('prompt_eval_ms_per_tok', '--'):>8.1f}" if isinstance(r.get('prompt_eval_ms_per_tok'), (int, float)) else f"{'--':>8}"
        gen = f"{r.get('gen_tok_s', '--'):>8.1f}" if isinstance(r.get('gen_tok_s'), (int, float)) else f"{'--':>8}"
        gen_ms = f"{r.get('gen_ms_per_tok', '--'):>8.1f}" if isinstance(r.get('gen_ms_per_tok'), (int, float)) else f"{'--':>8}"
        sg = f"{r.get('gen_sustained_tok_s', '--'):>8.1f}" if isinstance(r.get('gen_sustained_tok_s'), (int, float)) else f"{'--':>8}"
        name = r.get('name', '?')[:28]
        print(f"{name:<28} {pe:<12} {pe_ms:<12} {gen:<12} {gen_ms:<12} {sg:<12}")

    print(f"\n{'='*100}")


def main():
    results_file = "/tmp/benchmark_single_model_results.json"

    if os.path.exists(results_file):
        os.remove(results_file)

    print(f"SINGLE-MODEL FLAG BENCHMARK")
    print(f"Date: {time.strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"Model: {MODEL}")
    print(f"Binary: {BINARY}")
    print(f"Configs to test: {len(CONFIGS)}")
    print(f"{'='*70}\n")

    for config in CONFIGS:
        benchmark_config(config, results_file)
        time.sleep(2)

    print_summary(results_file)
    print(f"\nResults saved to {results_file}")


if __name__ == "__main__":
    main()