#!/usr/bin/env python3
"""Benchmark llama.cpp speculative configs on a fixed prompt set.

Protocol (see NGRAM_TASK.md ยง4): the same prompts, chat template and system
prompt across four configs, shuffled order, warmup excluded, acceptance taken
from /completion timings (draft_n / draft_n_accepted), pp reported to make
sure speculation does not hurt prompt processing.

Runs on the host that owns the model; starts and stops llama-server itself.

Usage (pilot):
    python3 ngram_bench.py --prompts ngram_prompts.json \
        --system-prompt ~/navi-system-prompt.txt --cache ngram-cache.bin \
        --configs A,B,C,D --ids short-ru-1,code-py-1,explain-ru-1,long-ru-1,short-en-1 \
        --reps 1 --temps 0 --n-predict 256 --out pilot.json
"""

from __future__ import annotations

import argparse
import glob
import http.client
import json
import os
import random
import re
import subprocess
import sys
import time
from urllib.parse import urlsplit

BIN = os.path.expanduser("~/llama/llama-b11960/llama-server")


def base_flags(ctx: int, ngl: int = 99, n_cpu_moe: int = 999,
               threads: int = 4) -> list[str]:
    """Flags shared by every config (the ones used in the earlier tg benches).

    ngl < 99 with n_cpu_moe 0 gives *whole-layer* offload: a layer lives
    entirely on one device, so its hidden state never crosses PCIe mid-layer.
    """
    flags = [
        "-c", str(ctx),
        "-ngl", str(ngl),
        "--n-cpu-moe", str(n_cpu_moe),
        "-fa", "on",
        "-ctk", "q8_0", "-ctv", "q8_0",
        "-t", str(threads),
        "--no-webui",
    ]
    return flags


def split_config(name: str) -> tuple[str, dict[str, int]]:
    """'B:32' -> config B with --n-cpu-moe 32 (experts only offload).
    'B@12' -> config B with -ngl 12 --n-cpu-moe 0 (whole layers on GPU).
    'B' keeps the defaults on every axis."""
    over: dict[str, int] = {}
    base = name
    if "@" in base:
        base, _, ngl = base.partition("@")
        over = {"ngl": int(ngl), "n_cpu_moe": 0}
    elif ":" in base:
        base, _, ncmoe = base.partition(":")
        over = {"n_cpu_moe": int(ncmoe)}
    return base, over


def config_flags(name: str, cache: str) -> list[str]:
    match name:
        case "A":
            return ["--spec-type", "none"]
        case "B":
            return ["--spec-type", "draft-mtp", "--spec-draft-n-max", "1"]
        case "C":
            return ["--spec-type", "ngram-cache", "-lcs", cache]
        case "D":
            return ["--spec-type", "draft-mtp,ngram-cache",
                    "--spec-draft-n-max", "1", "-lcs", cache]
        case "E":
            # Control: lookup decoding from the prompt's own context only.
            # C - E is the contribution of the corpus table.
            return ["--spec-type", "ngram-cache"]
        case "F":
            # Control: purely context-based ngram-mod, no corpus at all.
            return ["--spec-type", "ngram-mod"]
        case "G":
            # Control: MTP + lookup without a static table.
            return ["--spec-type", "draft-mtp,ngram-cache",
                    "--spec-draft-n-max", "1"]
        case "H":
            return ["--spec-type", "draft-mtp", "--spec-draft-n-max", "2"]
        case "I":
            return ["--spec-type", "draft-mtp", "--spec-draft-n-max", "3"]
        case "J":
            return ["--spec-type", "draft-mtp", "--spec-draft-n-max", "5"]
        case "K":
            return ["--spec-type", "draft-mtp,ngram-mod",
                    "--spec-draft-n-max", "1"]
        case "L":
            return ["--spec-type", "draft-mtp,ngram-map-k4v",
                    "--spec-draft-n-max", "1"]
        case "M":
            return ["--spec-type", "ngram-map-k4v,draft-mtp",
                    "--spec-draft-n-max", "1"]
        case _:
            raise SystemExit(f"unknown config {name!r}")


class Server:
    def __init__(self, model: str, port: int, flags: list[str], log_path: str,
                 ctx: int = 4096, ngl: int = 99, n_cpu_moe: int = 999,
                 threads: int = 4):
        self.port = port
        self.log_path = log_path
        self.log = open(log_path, "wb")
        self.proc = subprocess.Popen(
            [BIN, "-m", model, "--port", str(port), "-lv", "4",
             *base_flags(ctx, ngl, n_cpu_moe, threads), *flags],
            stdout=self.log, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL,
        )
        self.conn = http.client.HTTPConnection("127.0.0.1", port, timeout=900)

    def wait_ready(self, timeout: float = 600) -> None:
        deadline = time.time() + timeout
        while time.time() < deadline:
            if self.proc.poll() is not None:
                raise RuntimeError(f"server died, see {self.log_path}")
            try:
                c = http.client.HTTPConnection("127.0.0.1", self.port, timeout=5)
                c.request("GET", "/health")
                if c.getresponse().status == 200:
                    c.close()
                    return
                c.close()
            except OSError:
                pass
            time.sleep(1)
        raise RuntimeError(f"server not ready in {timeout}s")

    def post(self, path: str, payload: dict) -> dict:
        body = json.dumps(payload)
        self.conn.request("POST", path, body=body,
                          headers={"Content-Type": "application/json"})
        resp = self.conn.getresponse()
        raw = resp.read()
        if resp.status != 200:
            raise RuntimeError(f"{path} HTTP {resp.status}: {raw[:400]!r}")
        return json.loads(raw)

    def stop(self) -> None:
        try:
            self.proc.terminate()
            self.proc.wait(timeout=30)
        except subprocess.TimeoutExpired:
            self.proc.kill()
            self.proc.wait()
        self.log.close()

    def spec_statistics(self) -> dict[str, dict]:
        """Sum every 'statistics <type>:' line from the verbose log."""
        pattern = re.compile(
            r"statistics\s+([a-z0-9-]+):\s+"
            r"#calls\(b,g,a\)\s*=\s*(\d+)\s+(\d+)\s+(\d+),.*?"
            r"#gen drafts\s*=\s*(\d+),\s*#acc drafts\s*=\s*(\d+),\s*"
            r"#gen tokens\s*=\s*(\d+),\s*#acc tokens\s*=\s*(\d+)"
        )
        totals: dict[str, dict] = {}
        with open(self.log_path, "r", errors="replace") as f:
            for line in f:
                m = pattern.search(line)
                if not m:
                    continue
                name = m.group(1)
                # b, g, a, gen_drafts, acc_drafts, gen_tokens, acc_tokens
                vals = [int(x) for x in m.groups()[1:]]
                acc = totals.setdefault(name, dict.fromkeys(
                    ("calls", "gen_drafts", "acc_drafts", "gen_tokens",
                     "acc_tokens"), 0))
                acc["calls"] += vals[2]  # 'a' column: accepted calls
                acc["gen_drafts"] += vals[3]
                acc["acc_drafts"] += vals[4]
                acc["gen_tokens"] += vals[5]
                acc["acc_tokens"] += vals[6]
        return totals


def apply_template(server: Server, system_prompt: str, user_prompt: str,
                   tools: list | None = None) -> str:
    payload = {
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
    }
    if tools:
        payload["tools"] = tools
    data = server.post("/apply-template", payload)
    prompt = data.get("prompt")
    if not prompt:
        raise RuntimeError(f"/apply-template returned {data!r}")
    return prompt


def vram_used_mb() -> int | None:
    """amdgpud exposes this on every card; used to watch the expert offload."""
    for path in glob.glob("/sys/class/drm/card*/device/mem_info_vram_used"):
        try:
            with open(path) as f:
                return int(f.read()) // (1024 * 1024)
        except (OSError, ValueError):
            continue
    return None


def run_config(name, args, prompts, templated, order):
    base, over = split_config(name)
    flags = config_flags(base, args.cache)
    if args.extra:
        flags += args.extra.split()
    log_path = os.path.join(args.log_dir,
                            f"{args.tag}-{name.replace(':', '_').replace('@', 'ngl')}.log")
    print(f"\n=== config {name}: {' '.join(flags)} "
          f"[{over or 'default layout'}] ===\n  log: {log_path}",
          file=sys.stderr)
    server = Server(args.model, args.port, flags, log_path, ctx=args.ctx,
                    ngl=over.get("ngl", 99),
                    n_cpu_moe=over.get("n_cpu_moe", 999),
                    threads=args.threads)
    try:
        server.wait_ready()
        vram = vram_used_mb()
        print(f"  vram after load: {vram} MB", file=sys.stderr)
        # Warmup is excluded from the numbers; its stats are ignored because
        # acceptance is computed from per-request timings below.
        server.post("/completion", {
            "prompt": templated[prompts[0]["id"]],
            "n_predict": 16, "temperature": 0, "cache_prompt": False,
            "seed": 1,
        })

        records = []
        for rep in range(args.reps):
            for prompt_id in order:
                for temp in args.temps:
                    t0 = time.time()
                    out = server.post("/completion", {
                        "prompt": templated[prompt_id],
                        "n_predict": args.n_predict,
                        "temperature": temp,
                        "cache_prompt": False,
                        "seed": args.seed + rep,
                    })
                    tm = out.get("timings", {})
                    rec = {
                        "config": name,
                        "prompt_id": prompt_id,
                        "rep": rep,
                        "temp": temp,
                        "wall_s": round(time.time() - t0, 2),
                        "cache_n": tm.get("cache_n"),
                        "prompt_n": tm.get("prompt_n"),
                        "prompt_ms": tm.get("prompt_ms"),
                        "predicted_n": tm.get("predicted_n"),
                        "predicted_ms": tm.get("predicted_ms"),
                        "draft_n": tm.get("draft_n"),
                        "draft_n_accepted": tm.get("draft_n_accepted"),
                        # Kept for the self-echo audit: a real answer must not
                        # be a copy of the prompt.
                        "content_head": (out.get("content") or "")[:200],
                    }
                    records.append(rec)
                    print(
                        f"  {prompt_id:<14} rep{rep} t={temp}: "
                        f"pp {rec['prompt_n']} tok, "
                        f"tg {rec['predicted_n']} tok, "
                        f"{rec['predicted_n'] / rec['predicted_ms'] * 1000:5.2f} t/s, "
                        f"draft {rec['draft_n_accepted']}/{rec['draft_n']}",
                        file=sys.stderr,
                    )
    finally:
        server.stop()
    # The log is complete only after the process exits.
    return records, server.spec_statistics(), vram


def summarize(records, stats):
    tg_n = sum(r["predicted_n"] for r in records)
    tg_ms = sum(r["predicted_ms"] for r in records)
    pp_n = sum(r["prompt_n"] for r in records)
    pp_ms = sum(r["prompt_ms"] for r in records)
    draft_n = sum(r["draft_n"] or 0 for r in records)
    draft_acc = sum(r["draft_n_accepted"] or 0 for r in records)
    return {
        "requests": len(records),
        "tg_tps": round(tg_n / tg_ms * 1000, 2) if tg_ms else None,
        "pp_tps": round(pp_n / pp_ms * 1000, 2) if pp_ms else None,
        "tg_tokens": tg_n,
        "pp_tokens": pp_n,
        "draft_tokens": draft_n,
        "draft_accepted": draft_acc,
        "acceptance": round(draft_acc / draft_n, 3) if draft_n else None,
        "spec_stats": stats,
    }


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--prompts", required=True)
    ap.add_argument("--system-prompt", required=True)
    ap.add_argument("--cache", default="", help="static ngram cache for C and D")
    ap.add_argument("--model", default=os.path.expanduser(
        "~/models/Qwen3.6-35B-A3B-MTP-UD-IQ3_XXS.gguf"))
    ap.add_argument("--configs", default="A,B,C,D")
    ap.add_argument("--ids", default="", help="comma separated prompt ids (all if empty)")
    ap.add_argument("--reps", type=int, default=1)
    ap.add_argument("--temps", default="0", help="comma separated temperatures")
    ap.add_argument("--n-predict", type=int, default=256)
    ap.add_argument("--ctx", type=int, default=4096)
    ap.add_argument("--tools", default="",
                    help="JSON list of OpenAI-style tool schemas to render "
                         "into the chat template, as Navi does in production")
    ap.add_argument("--port", type=int, default=8299)
    ap.add_argument("--threads", type=int, default=4)
    ap.add_argument("--extra", default="",
                    help="extra llama-server flags, e.g. '--ub 2048 -b 4096'")
    ap.add_argument("--seed", type=int, default=1234)
    ap.add_argument("--tag", default="pilot")
    ap.add_argument("--log-dir", default=os.path.expanduser("~/ngram/smoke"))
    ap.add_argument("--out", default="")
    args = ap.parse_args()

    prompts = json.load(open(args.prompts, encoding="utf-8"))
    wanted = [p.strip() for p in args.ids.split(",") if p.strip()]
    if wanted:
        prompts = [p for p in prompts if p["id"] in wanted]
        missing = set(wanted) - {p["id"] for p in prompts}
        if missing:
            raise SystemExit(f"unknown prompt ids: {sorted(missing)}")
    args.temps = [float(t) for t in args.temps.split(",")]
    configs = [c.strip() for c in args.configs.split(",") if c.strip()]
    system_prompt = open(args.system_prompt, encoding="utf-8").read()
    tools = json.load(open(args.tools, encoding="utf-8")) if args.tools else None
    os.makedirs(args.log_dir, exist_ok=True)

    # One server just to render the chat template; the rendered text is
    # reused for every config so the prompt bytes are identical.
    tmpl_log = os.path.join(args.log_dir, f"{args.tag}-template.log")
    tmpl_server = Server(args.model, args.port, ["--spec-type", "none"],
                         tmpl_log, ctx=args.ctx)
    try:
        tmpl_server.wait_ready()
        templated = {
            p["id"]: apply_template(tmpl_server, system_prompt, p["prompt"], tools)
            for p in prompts
        }
    finally:
        tmpl_server.stop()

    # Report the prompt length once: it must leave room for n_predict inside ctx.
    sample = templated[prompts[0]["id"]]
    print(f"rendered prompt sample: {len(sample)} chars (~{len(sample)//3} tok)",
          file=sys.stderr)

    order = [p["id"] for p in prompts]
    results = {}
    for name in configs:
        random.Random(args.seed).shuffle(order)
        try:
            records, stats, vram = run_config(name, args, prompts,
                                              templated, order)
        except Exception as exc:  # an OOM config must not kill the whole sweep
            print(f"config {name} FAILED: {exc}", file=sys.stderr)
            results[name] = {"error": str(exc)}
            continue
        results[name] = {"summary": summarize(records, stats),
                         "vram_mb": vram, "records": records}

    print("\n=== summary ===")
    header = f"{'cfg':<6} {'req':>4} {'tg t/s':>8} {'pp t/s':>8} " \
             f"{'accept':>7} {'draft acc/gen':>14} {'vram MB':>8}"
    print(header)
    print("-" * len(header))
    for name, res in results.items():
        if "summary" not in res:
            print(f"{name:<6} FAILED: {res['error'][:60]}")
            continue
        s = res["summary"]
        print(f"{name:<6} {s['requests']:>4} {s['tg_tps']:>8} {s['pp_tps']:>8} "
              f"{(s['acceptance'] if s['acceptance'] is not None else '-'):>7} "
              f"{s['draft_accepted']:>6}/{s['draft_tokens']:<7} "
              f"{res['vram_mb'] if res['vram_mb'] is not None else '-':>8}")

    payload = {"args": vars(args), "results": results}
    if args.out:
        with open(args.out, "w", encoding="utf-8") as f:
            json.dump(payload, f, ensure_ascii=False, indent=2)
        print(f"\nwritten: {args.out}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())