#!/usr/bin/env python3
"""BFCL v4 tool-calling gate for BTL-4 Compact, runnable on a Colab T4.

BTL-3 Compact's weakest published number was parallel-multiple at 3/10, so tool
use is the measurement that matters most for an agent model and the one thing
the knowledge gate could not tell us. This runs the same 13 single-turn AST
categories, with the same system prompt and the same official `ast_checker`
that produced BTL-4's published 73.5% at bf16 -- so the result is directly
comparable rather than merely indicative.

The tool-call parser below is copied verbatim from btl_data_forge so that a
parsing difference cannot be mistaken for a quantisation effect.

Setup (one Colab cell, ~10 min for the CUDA build):

    !CMAKE_ARGS="-DGGML_CUDA=on" pip install -q llama-cpp-python
    !pip install -q bfcl-eval huggingface_hub
    !python bfcl_compact.py --cap 40

Full set is ~1240 cases and takes a couple of hours on a T4; --cap 40 samples
each category for a ~45 minute read. Report the cap you used -- a capped run is
not the full-set number and should never be printed as one.
"""
from __future__ import annotations

import argparse
import collections
import importlib.util
import json
import os
import time

# ---------------------------------------------------------------------------
# Verbatim copy of btl_data_forge.harness.toolparse -- format-tolerant parsing
# normalising every wrapper (bare JSON, <tool_call>, bracket DSL) to
# {"name": str, "arguments": dict}.
# ---------------------------------------------------------------------------

import ast
import json
import re

# key aliases seen across formats
_NAME_KEYS = ("name", "function_name", "function", "tool", "api")
_ARG_KEYS = ("arguments", "function_arg", "args", "parameters", "params", "arguments_json")


def _from_obj(obj) -> dict | None:
    if not isinstance(obj, dict):
        return None
    name = next((obj[k] for k in _NAME_KEYS if isinstance(obj.get(k), str)), None)
    if not name:
        return None
    args = next((obj[k] for k in _ARG_KEYS if k in obj), {})
    if isinstance(args, str):
        try:
            args = json.loads(args)
        except json.JSONDecodeError:
            args = {}
    if not isinstance(args, dict):
        args = {}
    return {"name": name, "arguments": args}


def _iter_json_objects(text: str):
    """Yield top-level {...} objects from text via brace-matching (handles nesting)."""
    depth = 0
    start = -1
    in_str = False
    esc = False
    for i, ch in enumerate(text):
        if in_str:
            if esc:
                esc = False
            elif ch == "\\":
                esc = True
            elif ch == '"':
                in_str = False
            continue
        if ch == '"':
            in_str = True
        elif ch == "{":
            if depth == 0:
                start = i
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0 and start >= 0:
                chunk = text[start : i + 1]
                try:
                    yield json.loads(chunk)
                except json.JSONDecodeError:
                    pass
                start = -1


def _parse_kwargs(argstr: str) -> dict:
    """Parse ToolACE bracket args: key="val", key=1, key=[1,2]. Robust to names with spaces."""
    out: dict = {}
    # split on top-level commas
    parts, depth, buf, in_str, q = [], 0, "", False, ""
    for ch in argstr:
        if in_str:
            buf += ch
            if ch == q:
                in_str = False
            continue
        if ch in "\"'":
            in_str, q = True, ch
            buf += ch
        elif ch in "([{":
            depth += 1
            buf += ch
        elif ch in ")]}":
            depth -= 1
            buf += ch
        elif ch == "," and depth == 0:
            parts.append(buf)
            buf = ""
        else:
            buf += ch
    if buf.strip():
        parts.append(buf)
    for p in parts:
        if "=" not in p:
            continue
        k, v = p.split("=", 1)
        k = k.strip()
        v = v.strip()
        try:
            out[k] = ast.literal_eval(v)
        except (ValueError, SyntaxError):
            out[k] = v.strip("\"'")
    return out


_CALL_HEAD = re.compile(r"([A-Za-z_][\w .-]*?)\s*\(")
_BARE_CALL_HEAD = re.compile(r"[A-Za-z_]\w*\s*\(")
_TOOLCALL_TAG = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.S)
_QWEN35_TOOLCALL = re.compile(
    r"<tool_call>\s*<function=([^>\n]+)>\s*(.*?)\s*</function>\s*</tool_call>",
    re.S,
)
_QWEN35_PARAMETER = re.compile(
    r"<parameter=([^>\n]+)>\s*(.*?)\s*</parameter>", re.S
)


def _parse_parameter_value(text: str):
    value = text.strip()
    try:
        return json.loads(value)
    except json.JSONDecodeError:
        try:
            return ast.literal_eval(value)
        except (ValueError, SyntaxError):
            return value


def _parse_qwen35_calls(text: str) -> list[dict]:
    calls: list[dict] = []
    for function_name, body in _QWEN35_TOOLCALL.findall(text):
        arguments = {
            name.strip(): _parse_parameter_value(value)
            for name, value in _QWEN35_PARAMETER.findall(body)
        }
        calls.append({"name": function_name.strip(), "arguments": arguments})
    return calls


def _find_calls(text: str) -> list[dict]:
    """Find Name(...) calls with balanced parens, so nested lists/dicts in args don't break it."""
    calls: list[dict] = []
    for m in _CALL_HEAD.finditer(text):
        name = m.group(1).strip()
        depth, j, in_str, q = 0, m.end() - 1, False, ""
        while j < len(text):
            ch = text[j]
            if in_str:
                if ch == q and text[j - 1] != "\\":
                    in_str = False
            elif ch in "\"'":
                in_str, q = True, ch
            elif ch == "(":
                depth += 1
            elif ch == ")":
                depth -= 1
                if depth == 0:
                    break
            j += 1
        if depth == 0 and j < len(text):
            argstr = text[m.end():j].strip()
            # some models emit a single positional dict: func({"a": 1, "b": 2})
            if argstr.startswith("{"):
                try:
                    d = json.loads(argstr)
                    if isinstance(d, dict):
                        calls.append({"name": name, "arguments": d})
                        continue
                except json.JSONDecodeError:
                    pass
            calls.append({"name": name, "arguments": _parse_kwargs(argstr)})
    return calls


def _find_bracket_calls(text: str) -> list[dict]:
    """Parse ToolACE calls only inside balanced square-bracket wrappers."""

    calls: list[dict] = []
    depth = 0
    start = -1
    in_str = False
    quote = ""
    escaped = False
    for index, char in enumerate(text):
        if in_str:
            if escaped:
                escaped = False
            elif char == "\\":
                escaped = True
            elif char == quote:
                in_str = False
            continue
        if char in "\"'":
            in_str = True
            quote = char
        elif char == "[":
            if depth == 0:
                start = index + 1
            depth += 1
        elif char == "]" and depth:
            depth -= 1
            if depth == 0 and start >= 0:
                calls.extend(_find_calls(text[start:index]))
                start = -1
    return calls


_THINK_END = re.compile(r"</think\s*>", re.I)


def _strip_thinking(text: str) -> str:
    """Parse only the answer, not the reasoning trace. Reasoning-model output is
    <think> ... </think> ANSWER; scanning the whole thing scrapes phantom calls out of the
    reasoning (e.g. "I could call foo(...)"), which wrongly fails abstention. Keep post-think text."""
    end = None
    for end in _THINK_END.finditer(text):
        pass
    return text[end.end():] if end else text


def parse_tool_calls(text: str) -> list[dict]:
    """Return a list of normalized {name, arguments} calls found in the model output."""
    text = _strip_thinking(text)
    calls: list[dict] = []

    # 1) Qwen3.5 native XML tool calls (strongest signal)
    calls = _parse_qwen35_calls(text)
    if calls:
        return calls

    # 2) Older Qwen JSON-in-<tool_call> tags
    tagged = _TOOLCALL_TAG.findall(text)
    if tagged:
        for chunk in tagged:
            try:
                c = _from_obj(json.loads(chunk))
                if c:
                    calls.append(c)
            except json.JSONDecodeError:
                pass
        if calls:
            return calls

    # 3) JSON objects anywhere (our model's function_name/function_arg style)
    for obj in _iter_json_objects(text):
        # a container like {"tool_calls":[...]} or {"function_calls":[...]} — dig in
        container = None
        if isinstance(obj, dict):
            for key in ("tool_calls", "function_calls"):
                if isinstance(obj.get(key), list):
                    container = obj[key]
                    break
        if container is not None:
            for sub in container:
                c = _from_obj(sub.get("function", sub) if isinstance(sub, dict) else sub)
                if c:
                    calls.append(c)
        else:
            c = _from_obj(obj)
            if c:
                calls.append(c)
    if calls:
        return calls

    # 4) ToolACE bracket DSL [Name(arg=val), ...]. Restrict the permissive
    # Name(...) parser to the bracket wrapper or a call at the very start of the
    # answer so prose parentheses cannot become hallucinated calls.
    calls = _find_bracket_calls(text)
    if calls:
        return calls
    stripped = text.lstrip()
    if _BARE_CALL_HEAD.match(stripped):
        return _find_calls(stripped)
    return []

# ---------------------------------------------------------------------------
# Gate
# ---------------------------------------------------------------------------

REPO = "badtheorylabs/BTL-4-Compact"
FILENAME = "BTL-4-IQ2_XXS.gguf"

SINGLE_TURN = [
    "simple_python", "simple_java", "simple_javascript", "multiple", "parallel",
    "parallel_multiple", "irrelevance", "live_simple", "live_multiple",
    "live_parallel", "live_parallel_multiple", "live_irrelevance",
    "live_relevance",
]

# Identical to bfcl_generate.py. Changing a word here breaks comparability with
# the published bf16 number, which is the entire point of this script.
SYS = ("You are an expert in composing functions. You are given a question and a set of possible "
       "functions. Based on the question, make one or more function/tool calls to achieve the "
       "purpose. If none of the functions can be used, point it out and do not call any function. "
       "Only return the function calls.\nHere is a list of functions in JSON format that you can "
       "invoke:\n")

NO_CALL = {"irrelevance", "live_irrelevance"}
YES_CALL = {"live_relevance"}


def data_dir() -> str:
    spec = importlib.util.find_spec("bfcl_eval")
    if spec is None or not spec.origin:
        raise SystemExit("pip install bfcl-eval")
    return os.path.join(os.path.dirname(spec.origin), "data")


def load_rows(ddir: str, cat: str, cap: int | None) -> list[dict]:
    path = os.path.join(ddir, f"BFCL_v4_{cat}.json")
    if not os.path.exists(path):
        return []
    rows = [json.loads(l) for l in open(path, encoding="utf-8") if l.strip()]
    return rows[:cap] if cap else rows


def build_messages(row: dict) -> list[dict]:
    msgs = [{"role": "system", "content": SYS + json.dumps(row["function"])}]
    for turn in row["question"][0]:
        msgs.append({"role": turn["role"], "content": turn["content"]})
    return msgs


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--cap", type=int, default=40,
                    help="cases per category; 0 = full set (~1240, hours on a T4)")
    ap.add_argument("--ctx", type=int, default=8192)
    ap.add_argument("--max-tokens", type=int, default=512)
    ap.add_argument("--out", default="bfcl_compact_cases.jsonl")
    ap.add_argument("--model", default=None, help="local .gguf path; default downloads from the Hub")
    args = ap.parse_args()
    cap = args.cap or None

    from llama_cpp import Llama

    path = args.model
    if path is None:
        from huggingface_hub import hf_hub_download
        print(f"fetching {REPO}/{FILENAME} ...", flush=True)
        path = hf_hub_download(repo_id=REPO, filename=FILENAME)

    print("loading (n_gpu_layers=-1 puts all 40 layers on the T4) ...", flush=True)
    llm = Llama(model_path=path, n_gpu_layers=-1, n_ctx=args.ctx,
                verbose=False, logits_all=False)

    ddir = data_dir()
    fout = open(args.out, "w", encoding="utf-8")
    t0, total = time.time(), 0

    for cat in SINGLE_TURN:
        rows = load_rows(ddir, cat, cap)
        if not rows:
            print(f"[skip] {cat}", flush=True)
            continue
        for row in rows:
            try:
                out = llm.create_chat_completion(
                    messages=build_messages(row), max_tokens=args.max_tokens,
                    temperature=0.0)
                raw = out["choices"][0]["message"].get("content") or ""
            except Exception as e:                       # keep going; record the miss
                raw = ""
                print(f"  !! {row['id']}: {type(e).__name__}", flush=True)
            parsed = [{c["name"]: c["arguments"]} for c in parse_tool_calls(raw)]
            fout.write(json.dumps({"cat": cat, "id": row["id"],
                                   "raw": raw[:4000], "parsed": parsed}) + "\n")
            total += 1
        fout.flush()
        print(f"[gen] {cat:<24} {len(rows):>4}  (total {total}, {time.time()-t0:.0f}s)",
              flush=True)
    fout.close()

    # ---- score with the official checker -----------------------------------
    import bfcl_eval
    from bfcl_eval.constants.enums import Language
    from bfcl_eval.constants.model_config import MODEL_CONFIG_MAPPING
    from bfcl_eval.eval_checker.ast_eval.ast_checker import ast_checker

    DATA = os.path.join(os.path.dirname(bfcl_eval.__file__), "data")
    PA = os.path.join(DATA, "possible_answer")
    LANG = {"simple_java": Language.JAVA, "simple_javascript": Language.JAVASCRIPT}
    literal = [m for m, c in MODEL_CONFIG_MAPPING.items()
               if not getattr(c, "underscore_to_dot", True)]
    MODEL_NAME = literal[0] if literal else next(iter(MODEL_CONFIG_MAPPING))

    def _load(d, cat, key):
        p = os.path.join(d, f"BFCL_v4_{cat}.json")
        if not os.path.exists(p):
            return {}
        return {json.loads(l)["id"]: json.loads(l)[key]
                for l in open(p, encoding="utf-8") if l.strip()}

    rows = [json.loads(l) for l in open(args.out, encoding="utf-8") if l.strip()]
    cats = sorted({r["cat"] for r in rows})
    funcs = {c: _load(DATA, c, "function") for c in cats}
    answers = {c: _load(PA, c, "ground_truth") for c in cats}

    tot: collections.Counter = collections.Counter()
    cor: collections.Counter = collections.Counter()
    for r in rows:
        cat, cid, parsed = r["cat"], r["id"], r["parsed"]
        tot[cat] += 1
        if cat in NO_CALL:
            cor[cat] += int(len(parsed) == 0)
            continue
        if cat in YES_CALL:
            cor[cat] += int(len(parsed) > 0)
            continue
        gt, fn = answers[cat].get(cid), funcs[cat].get(cid)
        if gt is None or fn is None:
            continue
        try:
            ok = bool(ast_checker(fn, parsed, gt, LANG.get(cat, Language.PYTHON),
                                  cat, MODEL_NAME).get("valid", False))
        except Exception:
            ok = False
        cor[cat] += int(ok)

    print(f"\n=== BFCL v4 single-turn — BTL-4 Compact (2.30 bpw) ===")
    print(f"{'category':<24}{'acc':>7}{'n':>8}")
    order = [c for c in SINGLE_TURN if c in tot]
    for c in order:
        print(f"{c:<24}{100*cor[c]/max(1,tot[c]):>6.1f}%{tot[c]:>8}")
    tc, tt = sum(cor.values()), sum(tot.values())
    micro = tc / max(1, tt)
    macro = sum(cor[c] / max(1, tot[c]) for c in order) / max(1, len(order))
    print(f"{'OVERALL (micro)':<24}{100*micro:>6.1f}%{tt:>8}")
    print(f"{'OVERALL (macro)':<24}{100*macro:>6.1f}%")
    print(f"\nbf16 reference (published, full set): 73.5% AST")
    if cap:
        print(f"NOTE: capped at {cap}/category — not the full-set number.")

    json.dump({"per_category": {c: {"acc": cor[c]/max(1,tot[c]), "n": tot[c]}
                                for c in order},
               "overall_micro": micro, "overall_macro": macro,
               "cap": cap, "bf16_reference_ast": 0.735},
              open("bfcl_compact_score.json", "w"), indent=2)
    print("wrote bfcl_compact_score.json")


if __name__ == "__main__":
    main()