#!/usr/bin/env python3
"""
Tool-selection benchmark for DiffusionGemma (zero-shot baseline or + LoRA adapter).

Generates with mlx-vlm's diffusion engine on the selector test set and scores
"- Tool" line predictions against expected tool sets: Jaccard, exact set match,
precision/recall, top-1. Numbers are comparable ONLY across runs of THIS harness
on the SAME test split (e.g. zero-shot vs adapter); cross-model or cross-harness
comparisons require re-running the other model through this same harness.
Note: per-sample outputs are order-coupled (device RNG is seeded once), so a
--max-samples k run replays the full run's first k samples; changing sample
order or count changes individual outputs.

Usage:
  python3 diffusion_eval.py \
      --model ./diffusiongemma-26B-A4B-it-4bit \
      [--adapter ./adapters/my-adapter] \
      --test ./data/test.jsonl --out ./eval.json \
      [--max-samples 50] [--check-template]
"""

import argparse
import json
import re
import time
from pathlib import Path

import mlx.core as mx

from mlx_vlm.utils import load as vlm_load

# Note: the model-turn prefill ("<|turn>model\n<|channel>thought\n<channel|>")
# is already baked into each test sample's prompt field by build_diffusiongemma_data.py,
# so generation starts exactly where training responses started.


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--model", required=True)
    p.add_argument("--adapter", default=None)
    p.add_argument("--test", required=True)
    p.add_argument("--out", required=True)
    p.add_argument("--max-samples", type=int, default=0, help="0 = all")
    p.add_argument("--max-tokens", type=int, default=96)
    p.add_argument("--sample-timeout", type=int, default=45,
                   help="per-sample wall-clock cap (s); DiffusionGemma's long-context prefill "
                        "can hang on ~900-token prompts — timed-out samples count as failures")
    p.add_argument("--seed", type=int, default=7)
    p.add_argument("--check-template", action="store_true",
                   help="verify our offline-rendered prompt matches the official template")
    return p.parse_args()


def parse_tools(text):
    """Extract '- Tool' lines; stop at turn end / eos markers."""
    tools = []
    for raw in text.split("\n"):
        line = raw.strip()
        for stop in ("<turn|>", "<eos>", "<|turn>"):
            if stop in line:
                line = line.split(stop)[0].strip()
        if line.startswith("- "):
            name = line[2:].strip()
            if name and re.fullmatch(r"[A-Za-z0-9_\-.]+", name):
                tools.append(name)
        elif tools and line and not line.startswith("-"):
            break  # left list format — stop collecting
    # de-dup preserving order
    seen, out = set(), []
    for t in tools:
        if t not in seen:
            seen.add(t)
            out.append(t)
    return out


def metrics(pred, expected):
    ps, es = set(pred), set(expected)
    inter = ps & es
    union = ps | es
    return {
        "jaccard": len(inter) / len(union) if union else 1.0,
        "exact": 1.0 if ps == es else 0.0,
        "precision": len(inter) / len(ps) if ps else 0.0,
        "recall": len(inter) / len(es) if es else 0.0,
        "top1": 1.0 if pred and pred[0] in es else 0.0,
    }


def main():
    args = parse_args()
    mx.random.seed(args.seed)

    print(f"[load] {args.model} adapter={args.adapter}", flush=True)
    model, processor = vlm_load(args.model, adapter_path=args.adapter)
    model.eval()
    tokenizer = processor.tokenizer
    mx.set_cache_limit(2 * 1024**3)  # bound the MLX buffer cache (paired with per-sample clear_cache)

    # resolve the generate API across mlx-vlm layouts
    try:
        from mlx_vlm import generate as vlm_generate
    except ImportError:
        from mlx_vlm.generate import generate as vlm_generate

    samples = []
    with open(args.test) as f:
        for line in f:
            obj = json.loads(line)
            expected_text = obj["response"].split("<turn|>")[0]
            expected = parse_tools(expected_text)
            samples.append({"prompt": obj["prompt"], "expected": expected})
    if args.max_samples:
        samples = samples[: args.max_samples]
    print(f"[data] {len(samples)} test samples", flush=True)

    # train/serve consistency guard: literal "<bos>" must tokenize to exactly one
    # leading id 2 (the generate path tokenizes with add_special_tokens=False)
    probe = tokenizer.encode("<bos>" + samples[0]["prompt"], add_special_tokens=False)
    assert probe[0] == 2 and 2 not in probe[1:], "double-BOS or missing BOS in eval tokenization"

    if args.check_template:
        from mlx_vlm.prompt_utils import apply_chat_template
        s = samples[0]["prompt"]
        # reconstruct messages from our rendered prompt
        sys_part = s.split("<|turn>system\n")[1].split("<turn|>")[0]
        usr_part = s.split("<|turn>user\n")[1].split("<turn|>")[0]
        messages = [{"role": "system", "content": sys_part},
                    {"role": "user", "content": usr_part}]
        official = apply_chat_template(processor, model.config, messages, add_generation_prompt=True)
        mine = s if s.startswith("<bos>") else "<bos>" + s
        print(f"[check] official == ours: {official == mine}", flush=True)
        if official != mine:
            print("OFFICIAL >>>", repr(official[-300:]), flush=True)
            print("OURS     >>>", repr(mine[-300:]), flush=True)

    # per-sample timeout: mlx-vlm's diffusion sampler is a Python loop calling mx
    # ops, so SIGALRM is delivered between denoising steps and aborts a wedged sample.
    import signal
    class _Timeout(Exception):
        pass
    def _on_alarm(signum, frame):
        raise _Timeout()
    signal.signal(signal.SIGALRM, _on_alarm)

    results, per_sample = [], []
    timeouts = 0
    t0 = time.time()
    for i, s in enumerate(samples):
        prompt_text = "<bos>" + s["prompt"]  # template normally injects bos; tokenized add_special_tokens=False
        timed_out = False
        try:
            signal.alarm(args.sample_timeout)
            out = vlm_generate(
                model, processor, prompt_text,
                max_tokens=args.max_tokens, temperature=0.0, verbose=False,
            )
            signal.alarm(0)
            text = out.text if hasattr(out, "text") else str(out)
        except _Timeout:
            signal.alarm(0)
            text = ""
            timed_out = True
            timeouts += 1
        pred = parse_tools(text)
        m = metrics(pred, s["expected"])
        results.append(m)
        per_sample.append({"i": i, "pred": pred, "expected": s["expected"],
                           "raw": text[:300], "timed_out": timed_out, **m})
        # CRITICAL: clear the MLX buffer cache between independent samples. Without
        # this the cache grows unboundedly over a long eval (RSS 16GB -> 22GB+),
        # progressively slowing generation until even short prompts hit the timeout.
        mx.clear_cache()
        if (i + 1) % 10 == 0:
            el = time.time() - t0
            print(f"  {i+1}/{len(samples)} jacc so far="
                  f"{sum(r['jaccard'] for r in results)/len(results):.4f} "
                  f"timeouts={timeouts} ({el/(i+1):.1f}s/sample)", flush=True)

    n = len(results)
    agg = {k: sum(r[k] for r in results) / n for k in results[0] if k != "timed_out"}
    agg["n"] = n
    agg["timeouts"] = timeouts
    agg["model"] = args.model
    agg["adapter"] = args.adapter
    agg["seconds_total"] = round(time.time() - t0, 1)

    print(json.dumps(agg, indent=2), flush=True)
    Path(args.out).write_text(json.dumps({"aggregate": agg, "samples": per_sample}, indent=2))
    print(f"[done] wrote {args.out}", flush=True)


if __name__ == "__main__":
    main()