#!/usr/bin/env python3
"""
Tool-selection benchmark for AUTOREGRESSIVE models via mlx_lm — the apples-to-apples
counterpart to diffusion_eval.py. Same clean test split, same parse_tools, same
metrics, so any AR model's number is directly comparable to DiffusionGemma's.

It reconstructs the raw system/user text from the DiffusionGemma-rendered prompt,
then re-renders it with the target model's OWN chat template (fair zero-shot framing).

Usage:
  python3 ar_eval.py \
      --model mlx-community/Qwen3.6-35B-A3B-4bit \
      --test ./data/test.jsonl --out ./eval_qwen.json \
      --label Qwen3.6-35B [--max-samples 0]
"""
import argparse, json, re, time
from pathlib import Path
import mlx.core as mx
from mlx_lm import load, generate
from mlx_lm.sample_utils import make_sampler


def parse_tools(text):
    tools = []
    for raw in text.split("\n"):
        line = raw.strip()
        for stop in ("<turn|>", "<eos>", "<|turn>", "</s>", "<|im_end|>", "<end_of_turn>"):
            if stop in line:
                line = line.split(stop)[0].strip()
        # accept "- Tool", "1. Tool", "* Tool", or bare "Tool" on its own line
        m = re.match(r"^[\-\*\d.\)\s]*([A-Za-z][A-Za-z0-9_\-.]+)\s*$", line)
        if line.startswith(("- ", "* ")) or re.match(r"^\d+[.)]", line):
            if m:
                tools.append(m.group(1))
    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, union = ps & es, ps | es
    return dict(
        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 reconstruct(prompt_text):
    """Pull raw system + user back out of the DiffusionGemma-rendered prompt."""
    sys_txt = prompt_text.split("<|turn>system\n", 1)[1].split("<turn|>", 1)[0].strip() \
        if "<|turn>system\n" in prompt_text else ""
    usr_txt = prompt_text.split("<|turn>user\n", 1)[1].split("<turn|>", 1)[0].strip()
    return sys_txt, usr_txt


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--test", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--label", required=True)
    ap.add_argument("--max-samples", type=int, default=0)
    ap.add_argument("--max-tokens", type=int, default=96)
    args = ap.parse_args()

    print(f"[load] {args.model}", flush=True)
    model, tokenizer = load(args.model)
    sampler = make_sampler(temp=0.0)  # greedy, deterministic

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

    results, per_sample = [], []
    t0 = time.time()
    for i, (sys_txt, usr_txt, expected) in enumerate(samples):
        msgs = ([{"role": "system", "content": sys_txt}] if sys_txt else []) + \
               [{"role": "user", "content": usr_txt}]
        try:
            prompt = tokenizer.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
        except Exception:
            prompt = (sys_txt + "\n\n" if sys_txt else "") + usr_txt
        text = generate(model, tokenizer, prompt=prompt, max_tokens=args.max_tokens,
                        sampler=sampler, verbose=False)
        pred = parse_tools(text)
        m = metrics(pred, expected)
        results.append(m)
        per_sample.append({"i": i, "pred": pred, "expected": expected, "raw": text[:300], **m})
        mx.clear_cache()  # bound MLX buffer cache over a long eval (see diffusion_eval.py)
        if (i + 1) % 20 == 0:
            print(f"  {i+1}/{len(samples)} jac so far="
                  f"{sum(r['jaccard'] for r in results)/len(results):.4f} "
                  f"({(time.time()-t0)/(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]}
    agg.update(n=n, model=args.model, label=args.label, 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] {args.out}", flush=True)


if __name__ == "__main__":
    main()