#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import subprocess
import time
from pathlib import Path

from harness_lib import SEEDS, Server, content, dump, message, peak_gpu

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


def run_bench(a, out: Path) -> None:
    env = os.environ.copy()
    env["LD_LIBRARY_PATH"] = f"/usr/local/cuda-11.8/lib64:{a.bench_bin.parent}:" + env.get("LD_LIBRARY_PATH", "")
    cmd = [
        str(a.bench_bin), "--model", str(a.model_path),
        "--n-prompt", "512,8192,32768", "--n-gen", "256",
        "--repetitions", "5", "--output", "json", "--progress",
        "--n-gpu-layers", "999", "--flash-attn", "on",
        "--cache-type-k", a.cache_k, "--cache-type-v", a.cache_v,
        "--batch-size", "2048", "--ubatch-size", "512", "--threads", "16",
    ]
    dump(out / "llama-bench-command.json", {"argv": cmd, "env": {"LD_LIBRARY_PATH": env["LD_LIBRARY_PATH"]}})
    gpu_path = out / "llama-bench-gpu.csv"
    with gpu_path.open("w") as gpu:
        monitor = subprocess.Popen(
            ["nvidia-smi", "--query-gpu=timestamp,memory.used,power.draw,utilization.gpu", "--format=csv,noheader,nounits", "-lms", "200"],
            stdout=gpu, stderr=subprocess.DEVNULL, text=True,
        )
        began = time.perf_counter()
        proc = subprocess.run(cmd, env=env, text=True, capture_output=True, timeout=3600)
        elapsed = time.perf_counter() - began
        monitor.terminate()
        try:
            monitor.wait(5)
        except subprocess.TimeoutExpired:
            monitor.kill()
    (out / "llama-bench.stderr.log").write_text(proc.stderr)
    (out / "llama-bench.stdout.json").write_text(proc.stdout)
    parsed = json.loads(proc.stdout) if proc.returncode == 0 else None
    dump(out / "llama-bench-summary.json", {"exit_code": proc.returncode, "elapsed_seconds": elapsed, "peak_gpu_mib": peak_gpu(gpu_path), "rows": parsed})


def build_prompt(server: Server, target_tokens: int) -> tuple[str, int]:
    unit = "Audit ledger cobalt vector maple quartz delta kernel harbor signal matrix. "
    unit_tokens = max(1, server.tokenize(unit) - 1)
    repeats = max(1, target_tokens // unit_tokens)
    text = unit * repeats
    count = server.tokenize(text)
    for _ in range(3):
        if abs(count - target_tokens) < max(64, target_tokens // 100):
            break
        repeats = max(1, int(repeats * target_tokens / max(1, count)))
        text = unit * repeats
        count = server.tokenize(text)
    return text, count


def context_probe(a, out: Path, ctx: int) -> dict:
    label = f"context-{a.cache_k}-{ctx}"
    server = Server(a.model_id, a.model_path, a.server_bin, out, ctx=ctx, reasoning="off", cache_k=a.cache_k, cache_v=a.cache_v, label=label)
    row = {"ctx_allocated": ctx, "cache_k": a.cache_k, "cache_v": a.cache_v, "passed": False}
    try:
        server.start()
        text, input_tokens = build_prompt(server, max(512, ctx - 2048))
        messages = [
            {"role": "system", "content": "Follow the exact response format."},
            {"role": "user", "content": text + "\nReply exactly TURN1_OK"},
        ]
        turns = []
        for number in [1, 2, 3]:
            run = server.complete(messages, seed=SEEDS[0], temperature=0, max_tokens=32)
            answer = content(run)
            expected = f"TURN{number}_OK"
            turns.append({"number": number, "passed": answer.strip() == expected, "output": answer, "elapsed_seconds": run["elapsed_seconds"], "usage": run["response"].get("usage"), "timings": run["response"].get("timings")})
            messages.append(message(run))
            if number < 3:
                messages.append({"role": "user", "content": f"Reply exactly TURN{number+1}_OK"})
        row.update({"passed": all(x["passed"] for x in turns), "input_tokens_tokenize": input_tokens, "input_characters": len(text), "turns": turns})
        first_timing = turns[0].get("timings") or {}
        prompt_ms = first_timing.get("prompt_ms")
        if prompt_ms:
            row["first_turn_characters_per_second"] = len(text) / (prompt_ms / 1000)
    except Exception as exc:
        row["error"] = str(exc)
    finally:
        row.update(server.stop())
    dump(out / f"{label}.json", row)
    return row


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--model-id", required=True)
    p.add_argument("--model-path", type=Path, required=True)
    p.add_argument("--root", type=Path, default=ROOT)
    p.add_argument("--server-bin", type=Path, default=ROOT / "llama.cpp/build/bin/llama-server")
    p.add_argument("--bench-bin", type=Path, default=ROOT / "llama.cpp/build/bin/llama-bench")
    p.add_argument("--cache-k", default="f16")
    p.add_argument("--cache-v", default="f16")
    p.add_argument("--contexts", default="8192,16384,32768,65536")
    p.add_argument("--skip-bench", action="store_true")
    a = p.parse_args()
    out = a.root / "results" / a.model_id / "performance" / f"{a.cache_k}-{a.cache_v}"
    out.mkdir(parents=True, exist_ok=True)
    if not a.skip_bench:
        print(json.dumps({"event":"bench_start","model":a.model_id}), flush=True)
        run_bench(a, out)
        print(json.dumps({"event":"bench_done","model":a.model_id}), flush=True)
    rows = []
    for ctx in [int(x) for x in a.contexts.split(",") if x.strip()]:
        print(json.dumps({"event":"context_start","model":a.model_id,"ctx":ctx,"cache":a.cache_k}), flush=True)
        rows.append(context_probe(a, out, ctx))
        print(json.dumps({"event":"context_done","model":a.model_id,"ctx":ctx,"passed":rows[-1]["passed"]}), flush=True)
    dump(out / "context-summary.json", {"rows": rows})


if __name__ == "__main__":
    main()

