#!/usr/bin/env python3
"""
Benchmark n-gram speculative decoding vs no-spec for frank-dragoon.

Tests prompt-processing throughput with 16K token prompts (manageable request size).
Code-heavy prompt (high n-gram redundancy) vs prose prompt (low redundancy).

Usage:
    # Generate prompts: python3 scripts/bench-ngram-mod.py generate /tmp/prompts/
    # Run: python3 scripts/bench-ngram-mod.py /tmp/prompts/
"""

import json
import os
import sys
import time
import subprocess
import urllib.request
import urllib.error
from pathlib import Path

MODEL = "/home/dank/models/kingjones-Qwen3.8-Flash-Next-ROCmFP4-STRIX/Qwen3.8-Flash-Next-Q4_0-ROCmFP4-STRIX.gguf"
CHAT_TEMPLATE = "/home/dank/serverAppConfigs/scripts/chat-template.jinja"
BUILD_DIR = "/home/dank/GitPackages/ROCmFPX/build-qwen38/bin"
LLAMA_BIN = os.path.join(BUILD_DIR, "llama-server")

# 16K tokens ~= 64KB text
TARGET_PROMPT_TOKENS = 16000
TOKEN_CHAR_RATIO = 4
TARGET_PROMPT_BYTES = TARGET_PROMPT_TOKENS * TOKEN_CHAR_RATIO

CONTEXT = 65536
BATCH = 4096
UBATCH = 4096
THREADS = 16
THREAD_BATCH = 32
NGRAM_N_MAX = 24
NGRAM_N_MIN = 4
TEST_TOKENS = 128
SERVER_WAIT = 300


def start_server(port, spec_type):
    cmd = [
        LLAMA_BIN,
        "--model", MODEL,
        "--chat-template-file", CHAT_TEMPLATE,
        "--jinja",
        "--device", "ROCm0",
        "--gpu-layers", "999",
        "--parallel", "1",
        "--ctx-size", str(CONTEXT),
        "--batch-size", str(BATCH),
        "--ubatch-size", str(UBATCH),
        "--fit", "off",
        "--flash-attn", "on",
        "--cache-type-k", "f16",
        "--cache-type-v", "f16",
        "--threads", str(THREADS),
        "--threads-batch", str(THREAD_BATCH),
        "--cache-prompt",
        "--cache-idle-slots",
        "--metrics",
        "--host", "127.0.0.1",
        "--port", str(port),
        "--temp", "0",
    ]
    if spec_type == "ngram-mod":
        cmd.extend([
            "--spec-type", "ngram-mod",
            "--spec-ngram-mod-n-max", str(NGRAM_N_MAX),
            "--spec-ngram-mod-n-min", str(NGRAM_N_MIN),
        ])
    else:
        cmd.extend(["--spec-type", "none"])

    env = os.environ.copy()
    env["ROCM_PATH"] = "/opt/rocm"
    env["LD_LIBRARY_PATH"] = f"{BUILD_DIR}:/opt/rocm/lib:/opt/rocm/lib/rocm_sysdeps/lib:/opt/rocm/lib/llvm/lib"
    env["HSA_OVERRIDE_GFX_VERSION"] = "11.5.1"
    env["HIP_VISIBLE_DEVICES"] = "0"

    proc = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    print(f"  Waiting for server on port {port} (spec={spec_type})...")
    for i in range(SERVER_WAIT):
        try:
            req = urllib.request.urlopen(f"http://127.0.0.1:{port}/health", timeout=2)
            resp = json.loads(req.read().decode())
            if resp.get("status") == "ok":
                print(f"  Server ready after {i+1}s")
                return proc
        except:
            time.sleep(1)

    proc.terminate()
    out, err = proc.communicate()
    print(f"ERROR: Server failed to start in {SERVER_WAIT}s")
    print(err.decode()[-2000:])
    sys.exit(1)


def run_test(prompt_text, spec_type, port):
    """Send a request and extract prompt processing metrics."""
    data = json.dumps({
        "model": "frank-dragoon",
        "messages": [{"role": "user", "content": prompt_text}],
        "max_tokens": TEST_TOKENS,
        "stream": False,
        "temperature": 0,
    }).encode()

    req = urllib.request.Request(
        f"http://127.0.0.1:{port}/v1/chat/completions",
        data=data,
        headers={"Content-Type": "application/json"},
    )

    start = time.monotonic()
    resp = urllib.request.urlopen(req, timeout=600)
    elapsed = time.monotonic() - start

    result = json.loads(resp.read().decode())
    timings = result.get("timings", {})
    usage = result.get("usage", {})

    return {
        "prompt_tokens": usage.get("prompt_tokens", 0),
        "prompt_ms": timings.get("prompt_ms", 0),
        "prompt_per_second": timings.get("prompt_per_second", 0),
        "prompt_per_token_ms": timings.get("prompt_per_token_ms", 0),
        "predicted_per_second": timings.get("predicted_per_second", 0),
        "total_time_s": elapsed,
        "cache_n": timings.get("cache_n", 0),
    }


def main():
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <prompt_dir>")
        sys.exit(1)

    prompt_dir = Path(sys.argv[1])
    prompt_dir.mkdir(parents=True, exist_ok=True)

    code_prompt_file = prompt_dir / "prompt_code.txt"
    prose_prompt_file = prompt_dir / "prompt_prose.txt"

    # Generate prompts if they don't exist
    if not code_prompt_file.exists() or not prose_prompt_file.exists():
        print("Generating prompts...")
        # Generate deterministic prompts
        # Code prompt: structured algorithms + data generation + config schemas
        code_lines = []
        for i in range(50):
            code_lines.append(f"# Algorithm {i}\n")
            code_lines.append(f"def algo_{i}(x, y):\n")
            code_lines.append(f"    \"\"\"Algorithm {i} processing.\"\"\"\n")
            code_lines.append(f"    result = 0\n")
            code_lines.append(f"    for j in range({i * 10}):\n")
            code_lines.append(f"        result += x * y + j\n")
            code_lines.append(f"    return result\n\n")
        code_text = "".join(code_lines)
        # Pad/truncate to target
        while len(code_text) < TARGET_PROMPT_BYTES:
            code_text += "# Additional algorithm block\n"
            code_text += f"def extra_{len(code_text)//100}():\n"
            code_text += "    x = 1\n"
            code_text += "    y = 2\n"
            code_text += "    return x + y\n\n"
        code_text = code_text[:TARGET_PROMPT_BYTES]

        # Prose prompt: historical/scientific/philosophical text
        prose_text = ""
        topics = [
            "The Renaissance brought profound transformation to European civilization. Artists and scholars pushed the boundaries of human expression.",
            "Quantum mechanics describes the behavior of matter and energy at atomic scales. The uncertainty principle limits measurement precision.",
            "Ethical philosophy grapples with questions of right and wrong. Utilitarianism judges actions by their consequences and the happiness they produce.",
            "Mountains rise like ancient sentinels, their peaks lost in clouds that swirl and dance across the sky above the valleys below.",
            "The ocean whispers secrets to the shore, each wave carrying fragments of distant lands across the water before retreating back to sea.",
        ]
        for _ in range(500):
            prose_text += f"{topics[len(prose_text) % len(topics)]} "
        prose_text = prose_text[:TARGET_PROMPT_BYTES]

        code_prompt_file.write_text(code_text)
        prose_prompt_file.write_text(prose_text)
        print(f"  Code: {code_prompt_file.stat().st_size:,} bytes")
        print(f"  Prose: {prose_prompt_file.stat().st_size:,} bytes")

    print("=" * 60)
    print("  N-gram Mod Benchmark — frank-dragoon")
    print(f"  STRIX ROCmFP4 · 64K context · 16K prompts")
    print("=" * 60)

    results = []
    port_base = 7772
    port_ngram = 7773

    tests = [
        ("code", "none", port_base, code_prompt_file),
        ("prose", "none", port_ngram, prose_prompt_file),
        ("code", "ngram-mod", port_base, code_prompt_file),
        ("prose", "ngram-mod", port_ngram, prose_prompt_file),
    ]

    for prompt_type, spec_type, port, pfile in tests:
        print(f"\n=== {spec_type} ({prompt_type}) ===")
        print(f"  File: {pfile} ({pfile.stat().st_size:,} bytes)")
        print(f"  Spec: {spec_type}")
        if spec_type == "ngram-mod":
            print(f"  n_max={NGRAM_N_MAX}, n_min={NGRAM_N_MIN}")

        proc = start_server(port, spec_type)

        try:
            # Warm-up
            print("  Warm-up...")
            run_test("Hello test", spec_type, port)

            # Full prompt
            print(f"  Sending {TARGET_PROMPT_TOKENS} token prompt...")
            with open(pfile) as f:
                prompt_text = f.read()
            result = run_test(prompt_text, spec_type, port)

            tps = result["prompt_per_second"]
            print(f"  Tokens:  {result['prompt_tokens']}")
            print(f"  Prompt:  {result['prompt_ms']:.0f}ms")
            print(f"  TPS:     {tps:.1f} tok/s")
            print(f"  Total:   {result['total_time_s']:.1f}s")

            results.append({
                "prompt_type": prompt_type,
                "spec_type": spec_type,
                "prompt_tokens": result["prompt_tokens"],
                "prompt_ms": result["prompt_ms"],
                "tps": tps,
                "total_time_s": result["total_time_s"],
            })

        finally:
            proc.terminate()
            try:
                proc.wait(timeout=5)
            except:
                proc.kill()

    # Summary
    print("\n" + "=" * 60)
    print("  RESULTS SUMMARY")
    print("=" * 60)

    for ptype in ["code", "prose"]:
        baseline = next((r for r in results if r["prompt_type"] == ptype and r["spec_type"] == "none"), None)
        ngram = next((r for r in results if r["prompt_type"] == ptype and r["spec_type"] == "ngram-mod"), None)
        if baseline and ngram:
            gain = ((ngram["tps"] - baseline["tps"]) / baseline["tps"]) * 100 if baseline["tps"] > 0 else 0
            print(f"\n{ptype.upper()} PROMPT:")
            print(f"  Baseline TPS: {baseline['tps']:.1f} tok/s ({baseline['prompt_ms']:.0f}ms)")
            print(f"  N-gram TPS:   {ngram['tps']:.1f} tok/s ({ngram['prompt_ms']:.0f}ms)")
            print(f"  Speedup:      {gain:+.1f}%")

    print("\n" + "=" * 60)


if __name__ == "__main__":
    main()