"""Record one model answering one prompt through llama-server, with a timestamp for every streamed chunk.

  python3 scripts/race_record.py --server BIN --model GGUF --name "Dory-Beam-4B" --loops 2 \
      --problem math500-123 --seed 0 --out race/dory-beam.json [--server-args "..."]

Starts llama-server, sends the problem as one user message (temperature 1.0, top_p 0.95, fixed seed), streams the
reply with raw <think> tags kept in the content, and writes {name, loops, problem, prompt, gold, events: [[t, text]],
completion_tokens, seconds, timings}. t is seconds since the request was sent. Stops the server afterwards.
"""

import argparse
import json
import shlex
import subprocess
import time
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
PORT = 8099


def wait_ready(proc, timeout=600):
    t0 = time.time()
    while time.time() - t0 < timeout:
        if proc.poll() is not None:
            raise RuntimeError("llama-server exited during startup")
        try:
            with urllib.request.urlopen(f"http://127.0.0.1:{PORT}/health", timeout=2) as r:
                if r.status == 200:
                    return
        except Exception:
            pass
        time.sleep(1)
    raise TimeoutError("llama-server did not become ready")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--server", required=True)
    ap.add_argument("--model", required=True)
    ap.add_argument("--name", required=True)
    ap.add_argument("--loops", type=int, default=2)
    ap.add_argument("--problem", required=True)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--max-tokens", type=int, default=16000)
    ap.add_argument("--server-args", default="")
    ap.add_argument("--effort", default="", help="reasoning_effort chat template argument (Dory-Beam-4B v2)")
    ap.add_argument("--out", required=True)
    args = ap.parse_args()
    prob = next(json.loads(l) for l in open(ROOT / "data" / "evalsets.jsonl") if json.loads(l)["id"] == args.problem)

    cmd = [args.server, "-m", args.model, "--host", "127.0.0.1", "--port", str(PORT), "-ngl", "999",
           "--ctx-size", "20480", "--parallel", "1", "--no-context-shift", "--jinja", "--reasoning-format", "none",
           "--override-kv", f"dory.num_loops=int:{args.loops},dory.recurrent_kv_cache_mode=str:per_loop",
           *shlex.split(args.server_args)]
    log = open(Path(args.out).with_suffix(".server.log"), "w")
    proc = subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT)
    try:
        wait_ready(proc)
        body = {"messages": [{"role": "user", "content": prob["prompt"]}], "stream": True, "temperature": 1.0,
                "top_p": 0.95, "top_k": 0, "min_p": 0, "seed": args.seed, "max_tokens": args.max_tokens,
                "cache_prompt": False, "timings_per_token": False}
        if args.effort:
            body["chat_template_kwargs"] = {"reasoning_effort": args.effort}
        req = urllib.request.Request(f"http://127.0.0.1:{PORT}/v1/chat/completions", json.dumps(body).encode(),
                                     {"Content-Type": "application/json"})
        events, timings, t0 = [], None, time.time()
        with urllib.request.urlopen(req, timeout=7200) as resp:
            for raw in resp:
                line = raw.decode("utf-8", "replace").strip()
                if not line.startswith("data: ") or line == "data: [DONE]":
                    continue
                chunk = json.loads(line[6:])
                timings = chunk.get("timings", timings)
                for ch in chunk.get("choices", []):
                    text = (ch.get("delta") or {}).get("content")
                    if text:
                        events.append([round(time.time() - t0, 3), text])
        seconds = time.time() - t0
    finally:
        proc.terminate()
        proc.wait(timeout=60)

    text = "".join(e[1] for e in events)
    out = {"name": args.name, "model_file": Path(args.model).name, "loops": args.loops, "effort": args.effort,
           "problem": args.problem,
           "prompt": prob["prompt"], "gold": prob["answer"], "seed": args.seed, "events": events,
           "completion_tokens": (timings or {}).get("predicted_n"), "seconds": round(seconds, 2), "timings": timings,
           "text": text}
    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    Path(args.out).write_text(json.dumps(out, ensure_ascii=False))
    print(f"{args.name}: {out['completion_tokens']} tokens in {seconds:.1f}s -> {args.out}")


if __name__ == "__main__":
    main()