#!/usr/bin/env python
"""Score a checkpoint (or the untouched base) on held-out agent tool-calling.

    python train/eval_tools.py --base Qwen/Qwen3.5-2B                      # baseline
    python train/eval_tools.py --base Qwen/Qwen3.5-2B --adapter runs/agent-2b/adapter

Reporting the base and the tuned model through the identical path is the whole
point — a tool-calling score means nothing without the lift it was measured
against.

Metrics, per task type and overall:
  tool_name_acc     did it call the tool the target called
  args_json_valid   did the emitted arguments parse at all
  args_exact        every argument name and value matches
  args_f1           per-field F1, so a near-miss scores better than a miss
  refusal_acc       on `refuse` rows: correctly declined to call anything
  price_acc         on `price` rows: picked the cheapest tool that fits

Qwen3.5 emits tool calls in an XML-ish form, not JSON:

    <tool_call>
    <function=NAME>
    <parameter=ARG>
    value
    </parameter>
    </function>
    </tool_call>

so `parse_tool_calls` reads that first and falls back to the Qwen2.5-style
JSON block, letting the same harness score either family.
"""

from __future__ import annotations

import argparse
import json
import re
import sys
import time
from collections import defaultdict
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO / "train"))

from data_utils import normalize_row  # noqa: E402

_XML_CALL = re.compile(r"<function=([^>]+)>(.*?)</function>", re.S)
_XML_PARAM = re.compile(r"<parameter=([^>]+)>\n?(.*?)\n?</parameter>", re.S)
_JSON_CALL = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.S)


def _coerce(text: str):
    """Parameters arrive as text; recover JSON values where they clearly are."""
    s = text.strip()
    if s and (s[0] in "{[" or s in ("true", "false", "null") or re.fullmatch(r"-?\d+(\.\d+)?", s)):
        try:
            return json.loads(s)
        except json.JSONDecodeError:
            return s
    return s


def parse_tool_calls(text: str) -> tuple[list[dict], bool]:
    """Return (calls, well_formed). `calls` is [{name, arguments}]."""
    calls: list[dict] = []
    ok = True

    for name, body in _XML_CALL.findall(text):
        args = {k.strip(): _coerce(v) for k, v in _XML_PARAM.findall(body)}
        calls.append({"name": name.strip(), "arguments": args})
    if calls:
        return calls, ok

    for blob in _JSON_CALL.findall(text):
        try:
            obj = json.loads(blob)
            args = obj.get("arguments", {})
            if isinstance(args, str):
                args = json.loads(args)
            calls.append({"name": obj.get("name", ""), "arguments": args if isinstance(args, dict) else {}})
        except json.JSONDecodeError:
            ok = False
    # A bare "<tool_call>" with nothing parseable is a malformed attempt, not a
    # clean refusal — distinguish so refusal accuracy stays honest.
    if not calls and "<tool_call>" in text:
        ok = False
    return calls, ok


def target_call(row: dict) -> dict | None:
    for msg in row["messages"]:
        if msg.get("role") == "assistant" and msg.get("tool_calls"):
            fn = msg["tool_calls"][0].get("function", msg["tool_calls"][0])
            args = fn.get("arguments")
            if isinstance(args, str):
                try:
                    args = json.loads(args)
                except json.JSONDecodeError:
                    args = {}
            return {"name": fn.get("name", ""), "arguments": args or {}}
    return None


def _extractability_of(raw: dict, gold: dict | None) -> dict | None:
    """Per-argument extractability, from the row if the generator recorded it,
    else recomputed from the user text. Returns None only when there is no gold
    call to grade."""
    if not gold:
        return None
    tagged = ((raw.get("meta") or {}).get("arg_extractable")
              if isinstance(raw.get("meta"), dict) else None)
    if isinstance(tagged, dict) and tagged:
        return tagged
    # Older rows predate the tag; recompute with the same rule so a corpus
    # rebuild is not required to get an honest number.
    user = next((m.get("content") or "" for m in raw.get("messages", [])
                 if m.get("role") == "user"), "")
    return {k: _is_extractable(v, user) for k, v in (gold.get("arguments") or {}).items()}


def _score_extractable(pred: dict, gold: dict, ext: dict | None) -> float:
    """1.0 when every EXTRACTABLE gold value is reproduced exactly and every
    SYNTHESISED gold key is present with a type-compatible value.

    The asymmetry is the point: reproducing a value the user stated is a real
    skill; reproducing one the dataset invented is impossible, so we ask only
    that the model supply *something of the right shape* for that field.
    """
    gold = gold or {}
    if not gold:
        return 1.0 if not pred else 0.0
    ext = ext or {}
    for k, gv in gold.items():
        if k not in (pred or {}):
            return 0.0
        pv = pred[k]
        if ext.get(k):
            if not _values_equal(pv, gv):
                return 0.0
        else:
            if not _type_compatible(pv, gv):
                return 0.0
    return 1.0


def _type_compatible(a, b) -> bool:
    """Same JSON shape. `5` vs `"5"` counts as compatible: an untyped schema
    accepts both and we do not want to fail a model for a defensible choice."""
    if isinstance(a, bool) or isinstance(b, bool):
        return isinstance(a, bool) and isinstance(b, bool)
    num = (int, float)
    if isinstance(b, num):
        if isinstance(a, num):
            return True
        return isinstance(a, str) and re.fullmatch(r"-?\d+(?:\.\d+)?", a.strip() or "") is not None
    if isinstance(b, str):
        return isinstance(a, (str, int, float))
    if isinstance(b, list):
        return isinstance(a, list)
    if isinstance(b, dict):
        return isinstance(a, dict)
    return type(a) is type(b)


def _values_equal(a, b) -> bool:
    if isinstance(a, (int, float)) and isinstance(b, (int, float)) \
            and not isinstance(a, bool) and not isinstance(b, bool):
        return abs(float(a) - float(b)) < 1e-9
    return str(a).strip().casefold() == str(b).strip().casefold()


def _is_extractable(value, user_text: str) -> bool:
    """Mirror of swarm-train/verify/reward.py's rule, inlined so the eval has no
    import dependency on the RL package. Kept deliberately identical."""
    if value is None or isinstance(value, bool):
        return False
    s = str(value).strip()
    if len(s) < 2:
        return False
    if s.casefold() in (user_text or "").casefold():
        return True
    if isinstance(value, (int, float)):
        for tok in re.findall(r"-?\d+(?:\.\d+)?", user_text or ""):
            try:
                if abs(float(tok) - float(value)) < 1e-9:
                    return True
            except ValueError:
                continue
    return False


def field_f1(pred: dict, gold: dict) -> float:
    if not pred and not gold:
        return 1.0
    pset = {(k, json.dumps(v, sort_keys=True)) for k, v in pred.items()}
    gset = {(k, json.dumps(v, sort_keys=True)) for k, v in gold.items()}
    if not pset or not gset:
        return 0.0
    tp = len(pset & gset)
    if tp == 0:
        return 0.0
    prec, rec = tp / len(pset), tp / len(gset)
    return 2 * prec * rec / (prec + rec)


def cheapest_tool(row: dict) -> str | None:
    """The lowest-priced tool offered, read from the price in its description."""
    best, best_price = None, None
    for tool in row.get("tools") or []:
        fn = tool.get("function", tool)
        m = re.search(r"\$\s*([0-9]*\.?[0-9]+)", fn.get("description", "") or "")
        if not m:
            continue
        price = float(m.group(1))
        if best_price is None or price < best_price:
            best, best_price = fn.get("name"), price
    return best


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", default="Qwen/Qwen3.5-2B")
    ap.add_argument("--adapter", default=None, help="LoRA adapter dir; omit to score the base model")
    ap.add_argument("--data", default=None, help="defaults to eval_ours.jsonl if present, else test.jsonl")
    ap.add_argument("--limit", type=int, default=None)
    ap.add_argument("--max-new-tokens", type=int, default=192)
    ap.add_argument("--out", default=None)
    args = ap.parse_args()

    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    data_dir = REPO / "data" / "out"
    path = Path(args.data) if args.data else (
        data_dir / "eval_ours.jsonl" if (data_dir / "eval_ours.jsonl").exists() else data_dir / "test.jsonl")
    rows = []
    with path.open() as fh:
        for line in fh:
            if line.strip():
                rows.append(json.loads(line))
    if args.limit:
        rows = rows[: args.limit]

    tok = AutoTokenizer.from_pretrained(args.base)
    model = AutoModelForCausalLM.from_pretrained(
        args.base,
        dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        device_map="auto" if torch.cuda.is_available() else None,
    )
    if args.adapter:
        from peft import PeftModel

        model = PeftModel.from_pretrained(model, args.adapter)
    model.eval()

    agg: dict[str, list] = defaultdict(list)
    per_task: dict[str, dict[str, list]] = defaultdict(lambda: defaultdict(list))
    t0 = time.time()

    for n, raw in enumerate(rows, 1):
        row = normalize_row(raw)
        task = row.get("task", "unknown")
        gold = target_call(row)

        # Feed everything up to the first assistant turn; the model produces it.
        cut = next((i for i, m in enumerate(row["messages"]) if m["role"] == "assistant"), len(row["messages"]))
        prompt = tok.apply_chat_template(
            row["messages"][:cut], tools=row.get("tools") or [],
            tokenize=False, add_generation_prompt=True)
        ids = tok(prompt, return_tensors="pt", add_special_tokens=False).to(model.device)
        with torch.no_grad():
            out = model.generate(**ids, max_new_tokens=args.max_new_tokens,
                                 do_sample=False, pad_token_id=tok.pad_token_id or tok.eos_token_id)
        text = tok.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True)
        calls, well_formed = parse_tool_calls(text)

        if task == "refuse":
            correct = len(calls) == 0
            agg["refusal_acc"].append(correct)
            per_task[task]["refusal_acc"].append(correct)
        elif gold:
            name_ok = bool(calls) and calls[0]["name"] == gold["name"]
            agg["tool_name_acc"].append(name_ok)
            per_task[task]["tool_name_acc"].append(name_ok)
            agg["args_json_valid"].append(well_formed and bool(calls))
            per_task[task]["args_json_valid"].append(well_formed and bool(calls))
            pred_args = calls[0]["arguments"] if calls else {}
            exact = name_ok and pred_args == gold["arguments"]
            f1 = field_f1(pred_args, gold["arguments"]) if name_ok else 0.0
            # `price` gold values are 0% extractable BY CONSTRUCTION — the task
            # is "pick the cheapest", and the arguments were invented wholesale
            # when the row was built. Reporting args_exact 0.0 there measures
            # the dataset, not the model; price_acc below is the real metric.
            report_args = _extractability_of(raw, gold) is not None and task != "price"
            if report_args:
                agg["args_exact"].append(exact)
                per_task[task]["args_exact"].append(exact)
                agg["args_f1"].append(f1)
                per_task[task]["args_f1"].append(f1)
                # THE HONEST METRIC: exact-match only values the user actually
                # stated; for values we synthesised, require the KEY to be
                # present and the TYPE to match, never the value. Grading an
                # unguessable value rewards fabrication.
                ex_ok = _score_extractable(pred_args, gold["arguments"],
                                           _extractability_of(raw, gold)) if name_ok else 0.0
                agg["args_exact_extractable"].append(ex_ok)
                per_task[task]["args_exact_extractable"].append(ex_ok)
            if task == "price":
                want = cheapest_tool(row) or gold["name"]
                hit = bool(calls) and calls[0]["name"] == want
                agg["price_acc"].append(hit)
                per_task[task]["price_acc"].append(hit)

        if n % 25 == 0:
            print(f"  {n}/{len(rows)}  ({time.time()-t0:.0f}s)", flush=True)

    def mean(xs):
        return round(sum(float(x) for x in xs) / len(xs), 4) if xs else None

    scorecard = {
        "base": args.base,
        "adapter": args.adapter,
        "data": str(path),
        "n": len(rows),
        "seconds": round(time.time() - t0, 1),
        "overall": {k: mean(v) for k, v in sorted(agg.items())},
        "by_task": {t: {k: mean(v) for k, v in sorted(m.items())} for t, m in sorted(per_task.items())},
    }

    out_path = Path(args.out) if args.out else (
        REPO / "runs" / f"eval-{'adapter' if args.adapter else 'base'}.json")
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(scorecard, indent=2))

    print(f"\n## {'tuned' if args.adapter else 'BASE'} — {len(rows)} rows, {scorecard['seconds']}s\n")
    print("| metric | score |\n|---|---:|")
    for k, v in scorecard["overall"].items():
        print(f"| {k} | {v} |")
    print("\n| task | " + " | ".join(sorted({k for m in per_task.values() for k in m})) + " |")
    keys = sorted({k for m in per_task.values() for k in m})
    print("|---" * (len(keys) + 1) + "|")
    for t, m in sorted(per_task.items()):
        print(f"| {t} | " + " | ".join(str(mean(m.get(k, []))) for k in keys) + " |")
    print(f"\nwrote {out_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())