#!/usr/bin/env python3
"""Separately labelled baseline-vs-MTP decode lane for the pinned llama.cpp build."""

from __future__ import annotations

import argparse
import json
import os
import statistics
import subprocess
import time
import urllib.request
from pathlib import Path

from harness_lib import dump, peak_gpu

ROOT = Path("/workspace/kingy-qwen38")
PORT = 8080


def request(path: str, payload: dict | None = None, timeout: int = 900):
    url = f"http://127.0.0.1:{PORT}{path}"
    if payload is None:
        req = urllib.request.Request(url)
    else:
        req = urllib.request.Request(
            url,
            data=json.dumps(payload).encode(),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
    with urllib.request.urlopen(req, timeout=timeout) as response:
        return json.load(response)


def summarize(rows: list[dict]) -> dict:
    ok = [row for row in rows if not row.get("error")]
    rates = [float(row["timings"]["predicted_per_second"]) for row in ok]
    elapsed = [float(row["elapsed_seconds"]) for row in ok]
    generated = sum(int(row["timings"].get("draft_n") or 0) for row in ok)
    accepted = sum(int(row["timings"].get("draft_n_accepted") or 0) for row in ok)
    return {
        "completed_runs": len(ok),
        "median_decode_tokens_per_second": statistics.median(rates) if rates else None,
        "median_elapsed_seconds": statistics.median(elapsed) if elapsed else None,
        "draft_tokens": generated,
        "draft_tokens_accepted": accepted,
        "acceptance_rate": accepted / generated if generated else None,
    }


def run_config(a, out: Path, label: str, extra: list[str]) -> dict:
    env = os.environ.copy()
    env["LD_LIBRARY_PATH"] = f"/usr/local/cuda-11.8/lib64:{a.server_bin.parent}:" + env.get("LD_LIBRARY_PATH", "")
    cmd = [
        str(a.server_bin), "--model", str(a.model_path), "--alias", a.model_id,
        "--host", "127.0.0.1", "--port", str(PORT), "--ctx-size", "8192",
        "--n-gpu-layers", "all", "--flash-attn", "on",
        "--cache-type-k", "f16", "--cache-type-v", "f16",
        "--batch-size", "2048", "--ubatch-size", "512", "--parallel", "1",
        "--threads", "16", "--threads-batch", "16", "--metrics", "--slots",
        "--reasoning", "off", "--jinja",
    ] + extra
    dump(out / f"{label}-command.json", {"argv": cmd, "env": {"LD_LIBRARY_PATH": env["LD_LIBRARY_PATH"]}})
    log_handle = (out / f"{label}.log").open("w")
    gpu_handle = (out / f"{label}-gpu.csv").open("w")
    monitor = subprocess.Popen(
        ["nvidia-smi", "--query-gpu=timestamp,memory.used,power.draw,utilization.gpu", "--format=csv,noheader,nounits", "-lms", "200"],
        stdout=gpu_handle,
        stderr=subprocess.DEVNULL,
        text=True,
    )
    began = time.perf_counter()
    proc = subprocess.Popen(cmd, env=env, stdout=log_handle, stderr=subprocess.STDOUT, text=True)
    result = {"label": label, "supported": False, "runs": []}
    try:
        deadline = time.time() + 900
        while time.time() < deadline:
            if proc.poll() is not None:
                raise RuntimeError((out / f"{label}.log").read_text(errors="replace")[-8000:])
            try:
                if request("/health", timeout=3).get("status") == "ok":
                    result["ready_seconds"] = time.perf_counter() - began
                    break
            except Exception:
                time.sleep(1)
        else:
            raise TimeoutError("server load timeout")
        payload = {
            "prompt": (
                "Audit ledger: cobalt=17, maple=29, quartz=41, delta=53. "
                "Continue with a detailed synthetic audit narrative. "
            ) * 18,
            "n_predict": 256,
            "temperature": 0,
            "seed": 20260816,
            "ignore_eos": True,
            "cache_prompt": False,
            "stream": False,
        }
        request("/completion", payload)
        for repetition in range(1, 6):
            start = time.perf_counter()
            response = request("/completion", payload)
            result["runs"].append({
                "repetition": repetition,
                "elapsed_seconds": time.perf_counter() - start,
                "timings": response.get("timings") or {},
                "tokens_predicted": response.get("tokens_predicted"),
                "stop_type": response.get("stop_type"),
                "response": response,
            })
        result["supported"] = True
        result["summary"] = summarize(result["runs"])
    except Exception as exc:
        result["error"] = str(exc)
    finally:
        if proc.poll() is None:
            proc.terminate()
            try:
                proc.wait(20)
            except subprocess.TimeoutExpired:
                proc.kill()
        if monitor.poll() is None:
            monitor.terminate()
            try:
                monitor.wait(5)
            except subprocess.TimeoutExpired:
                monitor.kill()
        log_handle.close()
        gpu_handle.close()
        result["peak_gpu_mib"] = peak_gpu(out / f"{label}-gpu.csv")
    dump(out / f"{label}-result.json", result)
    return result


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-id", required=True)
    parser.add_argument("--model-path", type=Path, required=True)
    parser.add_argument("--draft-model", type=Path)
    parser.add_argument("--root", type=Path, default=ROOT)
    parser.add_argument("--server-bin", type=Path, default=ROOT / "llama.cpp/build/bin/llama-server")
    a = parser.parse_args()
    out = a.root / "results" / a.model_id / "mtp"
    out.mkdir(parents=True, exist_ok=True)
    baseline = run_config(a, out, "baseline", [])
    mtp_args = ["--spec-type", "draft-mtp", "--spec-draft-ngl", "all", "--spec-draft-n-max", "3"]
    if a.draft_model:
        mtp_args += ["--spec-draft-model", str(a.draft_model)]
    mtp = run_config(a, out, "mtp", mtp_args)
    speedup = None
    extra_vram = None
    if baseline.get("supported") and mtp.get("supported"):
        base_rate = baseline["summary"]["median_decode_tokens_per_second"]
        mtp_rate = mtp["summary"]["median_decode_tokens_per_second"]
        speedup = mtp_rate / base_rate if base_rate else None
        if baseline.get("peak_gpu_mib") is not None and mtp.get("peak_gpu_mib") is not None:
            extra_vram = mtp["peak_gpu_mib"] - baseline["peak_gpu_mib"]
    dump(out / "summary.json", {
        "model_id": a.model_id,
        "draft_model": str(a.draft_model) if a.draft_model else None,
        "baseline": baseline,
        "mtp": mtp,
        "mtp_speedup_ratio": speedup,
        "mtp_extra_peak_vram_mib": extra_vram,
    })
    print(json.dumps({"model": a.model_id, "baseline_supported": baseline.get("supported"), "mtp_supported": mtp.get("supported"), "speedup": speedup, "extra_vram_mib": extra_vram}), flush=True)


if __name__ == "__main__":
    main()
