#!/usr/bin/env python3
"""Diverse, paired practical-output validation for MTP speculative decoding."""

from __future__ import annotations

import argparse
import hashlib
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
PROMPTS = [
    {
        "id": "technical_memo",
        "text": "Write a 180–220 word technical memo explaining why a Q4 model that loads at 8K context may still fail at 64K. Discuss weights, KV cache, runtime allocations, output reserve, and measured headroom. Use plain English and do not invent benchmark numbers.",
    },
    {
        "id": "coding_review",
        "text": "Write a 180–220 word code-review response. A Python function intended to merge overlapping or touching integer intervals currently returns sorted(items). Explain the defect, provide a corrected function, cover empty input and input immutability, and name three tests.",
    },
    {
        "id": "incident_runbook",
        "text": "Create exactly 12 numbered incident-response steps for a local LLM server that suddenly returns malformed tool calls. Cover reproduction, frozen fixtures, chat template, parser, runtime revision, model hash, logs, rollback, and verification. Each step must be one concise sentence.",
    },
]


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 sha(text: str) -> str:
    return hashlib.sha256(text.encode()).hexdigest()


def summarize(rows: list[dict]) -> dict:
    by_prompt = {}
    for prompt in PROMPTS:
        subset = [row for row in rows if row["prompt_id"] == prompt["id"] and not row.get("error")]
        rates = [float(row["timings"]["predicted_per_second"]) for row in subset]
        elapsed = [float(row["elapsed_seconds"]) for row in subset]
        by_prompt[prompt["id"]] = {
            "completed_runs": len(subset),
            "median_decode_tokens_per_second": statistics.median(rates) if rates else None,
            "median_end_to_end_seconds": statistics.median(elapsed) if elapsed else None,
            "median_completion_tokens": statistics.median([row["completion_tokens"] for row in subset]) if subset else None,
        }
    draft_n = sum(int(row["timings"].get("draft_n") or 0) for row in rows if not row.get("error"))
    accepted = sum(int(row["timings"].get("draft_n_accepted") or 0) for row in rows if not row.get("error"))
    return {
        "completed_runs": sum(not row.get("error") for row in rows),
        "by_prompt": by_prompt,
        "draft_tokens": draft_n,
        "draft_tokens_accepted": accepted,
        "acceptance_rate": accepted / draft_n if draft_n 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,
    )
    proc = subprocess.Popen(cmd, env=env, stdout=log_handle, stderr=subprocess.STDOUT, text=True)
    result = {"label": label, "supported": False, "runs": []}
    began = time.perf_counter()
    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")
        warm = {"model": a.model_id, "messages": [{"role": "user", "content": "Reply with one sentence about local inference."}], "temperature": 0, "seed": 20260816, "max_tokens": 96, "stream": False}
        request("/v1/chat/completions", warm)
        for prompt in PROMPTS:
            payload = {"model": a.model_id, "messages": [{"role": "user", "content": prompt["text"]}], "temperature": 0, "seed": 20260816, "max_tokens": 384, "stream": False}
            for repetition in range(1, 6):
                started = time.perf_counter()
                response = request("/v1/chat/completions", payload)
                elapsed = time.perf_counter() - started
                output = response["choices"][0]["message"].get("content") or ""
                timings = response.get("timings") or {}
                result["runs"].append({
                    "prompt_id": prompt["id"], "repetition": repetition,
                    "elapsed_seconds": elapsed, "completion_tokens": int((response.get("usage") or {}).get("completion_tokens") or timings.get("predicted_n") or 0),
                    "output": output, "output_sha256": sha(output), "timings": timings,
                })
        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("--disable-draft-backend-sampling", action="store_true")
    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()
    lane = "mtp-practical-cpu-draft-sampling" if a.disable_draft_backend_sampling else "mtp-practical"
    out = a.root / "results" / a.model_id / lane
    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.disable_draft_backend_sampling: mtp_args += ["--no-spec-draft-backend-sampling"]
    if a.draft_model: mtp_args += ["--spec-draft-model", str(a.draft_model)]
    mtp = run_config(a, out, "mtp", mtp_args)
    paired = []
    if baseline.get("supported") and mtp.get("supported"):
        for left, right in zip(baseline["runs"], mtp["runs"]):
            paired.append({
                "prompt_id": left["prompt_id"], "repetition": left["repetition"],
                "outputs_identical": left["output_sha256"] == right["output_sha256"],
                "end_to_end_speedup_ratio": left["elapsed_seconds"] / right["elapsed_seconds"] if right["elapsed_seconds"] else None,
                "decode_speedup_ratio": right["timings"].get("predicted_per_second", 0) / left["timings"].get("predicted_per_second", 1),
            })
    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, "paired": paired,
        "identical_outputs": sum(row["outputs_identical"] for row in paired), "paired_runs": len(paired),
        "median_end_to_end_speedup_ratio": statistics.median([row["end_to_end_speedup_ratio"] for row in paired]) if paired else None,
        "median_decode_speedup_ratio": statistics.median([row["decode_speedup_ratio"] for row in paired]) if paired else None,
        "mtp_extra_peak_vram_mib": (mtp.get("peak_gpu_mib") - baseline.get("peak_gpu_mib")) if paired else None,
    })
    print(json.dumps({"model": a.model_id, "baseline_supported": baseline.get("supported"), "mtp_supported": mtp.get("supported"), "paired_runs": len(paired), "identical": sum(row["outputs_identical"] for row in paired), "median_e2e_speedup": statistics.median([row["end_to_end_speedup_ratio"] for row in paired]) if paired else None}), flush=True)


if __name__ == "__main__":
    main()
