#!/usr/bin/env python3
"""voice-check: one-command capability check for a served model.

Assumes a llama-server with your model on --upstream (default
http://localhost:8080). Runs ARC (100), MMLU (2 x 57 subjects) and GSM8K
through lm-eval using this repo's benchmarking venv, extracts true scores
with the thinking-model-aware rescore, and writes everything into
<out>/<label>/:

  scorecard.md        the numbers table (also printed)
  samples/            every prompt, answer, and score (harness JSONL)
  thinking.jsonl      every prompt, answer, and reasoning trace

Needs: benchmarking/.venv with lm-eval (see benchmarking/requirements.txt).
Temperature 0.0, seed 1234, budgets ARC 1024 / MMLU 2048 / GSM8K 1024 —
identical for every model so deltas stay comparable. Unanswered items
(empty content) are counted separately, never as wrong.

Usage:
  python3 voice-check.py --label file-c [--upstream URL] [--out DIR]
  python3 voice-check.py --label smoke --smoke   # 1 ARC item, ~1 min
"""
import argparse
import os
import re
import shutil
import socket
import subprocess
import sys
import time
import urllib.request
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
VENV_PY = os.path.join(ROOT, "benchmarking", ".venv", "bin", "python")
SHIM = os.path.join(ROOT, "benchmarking", "shim.py")
RESCORE = os.path.join(ROOT, "benchmarking", "rescore.py")

TASKS = [  # (name, task, limit, budget, until_override); until matches the
    # published runs exactly: only MMLU overrides task stops
    ("arc", "arc_challenge_chat", 100, 1024, False),
    ("mmlu", "mmlu_generative", 2, 2048, True),
    ("gsm8k", "tinyGSM8k", None, 1024, False),
]


def free_port():
    s = socket.socket()
    s.bind(("127.0.0.1", 0))
    p = s.getsockname()[1]
    s.close()
    return p


def wait_http(url, timeout=60):
    t0 = time.time()
    while time.time() - t0 < timeout:
        try:
            urllib.request.urlopen(url, timeout=5).read(1)
            return True
        except Exception:
            time.sleep(1)
    return False


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--label", required=True)
    ap.add_argument("--upstream", default="http://localhost:8080")
    ap.add_argument("--out", default=".")
    ap.add_argument("--smoke", action="store_true")
    ap.add_argument("--force", action="store_true",
                    help="overwrite an existing output folder")
    a = ap.parse_args()

    if not os.path.isfile(VENV_PY):
        sys.exit("benchmarking venv missing: uv venv benchmarking/.venv && "
                 "uv pip install --python benchmarking/.venv/bin/python "
                 "-r benchmarking/requirements.txt")
    if not a.label or a.label.strip(" .") != a.label or any(
            c in a.label for c in ("/", "\\", "..")):
        sys.exit("--label must be a plain folder name")
    try:
        with urllib.request.urlopen(a.upstream.rstrip("/") + "/v1/models",
                                     timeout=15) as r:
            r.read(1)
    except Exception:
        sys.exit(f"no model server at {a.upstream} "
                 "(serve it first, then rerun)")
    dest = os.path.join(a.out, a.label)
    if os.path.exists(os.path.join(dest, "scorecard.md")) and not a.force:
        sys.exit(f"{dest}/scorecard.md exists: rerun with --force to redo")
    os.makedirs(os.path.join(dest, "samples"), exist_ok=True)

    port = free_port()
    shim = subprocess.Popen(
        [sys.executable, SHIM, a.upstream.rstrip("/"), str(port)],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    think_dest = os.path.join(dest, "thinking.jsonl")
    think_src = os.path.join(ROOT, "benchmarking", "thinking-log.jsonl")
    think_mark = os.path.getsize(think_src) if os.path.exists(think_src) else 0
    try:
        if not wait_http(f"http://127.0.0.1:{port}/tokenizer_info"):
            sys.exit("shim failed to start")
        base = f"http://127.0.0.1:{port}"
        rows = {}
        tasks = ([("arc", "arc_challenge_chat", 1, 1024, False)] if a.smoke
                 else TASKS)
        for name, task, limit, budget, over in tasks:
            gen = (f"max_gen_toks={budget},until=[\"</s>\"]" if over
                   else f"max_gen_toks={budget}")
            cmd = [VENV_PY, "-m", "lm_eval", "run",
                   "--model", "local-chat-completions",
                   "--model_args",
                   f"base_url={base}/v1/chat/completions,model={a.label},"
                   f"num_concurrent=2,max_retries=3",
                   "--tasks", task, "--apply_chat_template",
                   "--gen_kwargs", gen,
                   "--log_samples",
                   "--output_path", os.path.join(dest, "raw", name)]
            if limit:
                cmd += ["--limit", str(limit)]
            print("benchmarking your model..", name, flush=True)
            r = subprocess.run(cmd)
            if r.returncode != 0:
                sys.exit(f"task {name} failed")
            kind = "gsm8k" if name == "gsm8k" else "letter"
            agg = [0, 0, 0]
            for f in sorted(os.listdir(
                    os.path.join(dest, "raw", name, a.label))):
                if not f.startswith("samples_"):
                    continue
                src = os.path.join(dest, "raw", name, a.label, f)
                p = subprocess.run(
                    [sys.executable, RESCORE, src, kind],
                    capture_output=True, text=True)
                m = re.search(r"n=(\d+) correct=(\d+) wrong=(\d+) "
                              r"unanswered=(\d+)", p.stdout)
                if not m:
                    sys.exit(f"rescore failed on {src}:\n{p.stdout}"
                             f"{p.stderr}")
                import itertools as _it
                for i, v in enumerate(int(x) for x in m.groups()[1:]):
                    agg[i] += v
                shutil.copy(src, os.path.join(
                    dest, "samples", f"{name}-{f}"))
            rows[name] = agg
        lines = ["# Scorecard: " + a.label, "",
                 "| Task | n | Correct | Wrong | Unanswered | Acc |",
                 "|---|---|---|---|---|---|"]
        for name, (c, w, u) in rows.items():
            n = c + w + u
            lines.append(f"| {name.upper()} | {n} | {c} | {w} | {u} | "
                         f"{c / n:.3f}" if n else f"| {name.upper()} | 0 | "
                         f"0 | 0 | 0 | n/a |")
        lines += ["",
                  "Scored with thinking-model-aware extraction "
                  "(benchmarking/rescore.py rules). Unanswered = empty "
                  "content, counted separately. Full prompts, answers, and "
                  "reasoning traces in samples/ and thinking.jsonl."]
        open(os.path.join(dest, "scorecard.md"), "w").write("\n".join(lines)
                                                            + "\n")
        print("\n".join(lines))
    finally:
        shim.terminate()
    # slice only this run's thinking entries out of the shared log
    if os.path.exists(think_src):
        with open(think_src, "rb") as f:
            f.seek(think_mark)
            tail = f.read()
        with open(think_dest, "wb") as f:
            f.write(tail)
    shutil.rmtree(os.path.join(dest, "raw"), ignore_errors=True)
    print(f"saved to {dest}/")

if __name__ == "__main__":
    main()