"""BFCL v4 loading + synthetic per-user session construction.

BFCL examples are almost all distinct functions, so
there is essentially no natural cross-session repetition. To test the
personalization/persistence claim we synthesize realistic repetition: each
simulated user is assigned a small set of *signature tasks* (real BFCL
examples), and across ordered sessions they re-issue those tasks with perturbed
numeric argument values -- i.e. the same tool used again with new inputs, which
is exactly the recurring per-user tool usage the memory is meant to exploit.

Crucially, we do NOT fabricate the target tool call: every (possibly perturbed)
query is sent to the genuinely served gpt-oss-120b model, and the model's real
generation is the target. The perturbation only changes the natural-language
query; the model decides what to emit.
"""
from __future__ import annotations

import json
import random
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

DATA_DIR = Path(__file__).resolve().parent.parent / "data" / "bfcl"
SEALTOOLS_DIR = Path(__file__).resolve().parent.parent / "data" / "sealtools"
TAU2_DIR = Path(__file__).resolve().parent.parent / "data" / "tau2"
TOOLALPACA_DIR = Path(__file__).resolve().parent.parent / "data" / "toolalpaca"
APIBANK_DIR = Path(__file__).resolve().parent.parent / "data" / "apibank"
TOOLBENCH_DIR = Path(__file__).resolve().parent.parent / "data" / "toolbench"


def _toolbench_schema(api: dict) -> dict[str, Any]:
    props, required = {}, []
    for p in (api.get("required_parameters") or []):
        props[p["name"]] = {"type": _SEAL_TYPE_MAP.get(str(p.get("type", "str")).lower(),
                                                        "string"),
                            "description": str(p.get("description", ""))[:160]}
        if props[p["name"]]["type"] not in ("string", "integer", "number",
                                            "boolean", "array", "object"):
            props[p["name"]]["type"] = "string"
        required.append(p["name"])
    for p in (api.get("optional_parameters") or []):
        t = _SEAL_TYPE_MAP.get(str(p.get("type", "str")).lower(), "string")
        if t not in ("string", "integer", "number", "boolean", "array", "object"):
            t = "string"
        props[p["name"]] = {"type": t,
                            "description": str(p.get("description", ""))[:160]}
    name = str(api.get("api_name", "api")).strip().replace(" ", "_")
    return {"name": name or "api",
            "description": str(api.get("api_description", ""))[:300],
            "parameters": {"type": "dict", "properties": props,
                           "required": required}}


def load_toolbench(splits: tuple[str, ...] = ("G1_instruction", "G2_instruction",
                                             "G3_instruction")) -> list[Task]:
    """ToolBench (Qin et al., 2023, arXiv:2307.16789) as Tasks, from the REAL
    dataset (OpenBMB/ToolBench Google-Drive `data/test_instruction/`). Each item
    has a natural-language `query` and an `api_list` of candidate tools with
    required/optional parameter schemas."""
    tasks: list[Task] = []
    for split in splits:
        f = TOOLBENCH_DIR / "test_instruction" / f"{split}.json"
        if not f.exists():
            continue
        for r in json.loads(f.read_text()):
            q = (r.get("query") or "").strip()
            apis = r.get("api_list") or []
            if not q or len(q) < 8 or not apis:
                continue
            funcs = [_toolbench_schema(a) for a in apis]
            # drop dup / empty-name schemas
            seen, uniq = set(), []
            for fn in funcs:
                if fn["name"] and fn["name"] not in seen:
                    seen.add(fn["name"])
                    uniq.append(fn)
            if not uniq:
                continue
            tid = f"toolbench_{split}_{r.get('query_id')}"
            tasks.append(Task(id=tid, query=q, functions=uniq, origin_id=tid))
    return tasks


def _toolalpaca_params(desc: str) -> dict[str, Any]:
    """Parse the 'Parameters: {...}' JSON-ish blob out of a ToolAlpaca function
    description into BFCL-style {type, description} properties. Best-effort:
    ToolAlpaca param values are free text ('string. One of: [...]'), so we keep
    the leading type word and stash the rest as the description."""
    m = re.search(r"Parameters:\s*(\{.*?\})\s*(?:\nOutput|$)", desc, re.S)
    props: dict[str, Any] = {}
    if not m:
        return props
    try:
        raw = json.loads(m.group(1))
    except Exception:
        return props
    for pname, pdesc in raw.items():
        head = str(pdesc).split(".")[0].strip().lower()
        typ = head if head in ("string", "integer", "number", "boolean",
                               "array", "object") else "string"
        prop: dict[str, Any] = {"type": typ, "description": str(pdesc)[:200]}
        if typ == "array":
            prop["items"] = {"type": "string"}
        props[pname] = prop
    return props


def load_toolalpaca(splits: tuple[str, ...] = ("eval_simulated",
                                              "eval_real")) -> list[Task]:
    """ToolAlpaca (Tang et al., 2023) as Tasks. Each API's Instructions become
    queries; the API's Function_Description entries become the tool registry
    for that query. Downloaded from github.com/tangqiaoyu/ToolAlpaca."""
    tasks: list[Task] = []
    for split in splits:
        f = TOOLALPACA_DIR / f"{split}.json"
        if not f.exists():
            continue
        for tool in json.loads(f.read_text()):
            fdesc = tool.get("Function_Description") or {}
            if not isinstance(fdesc, dict):
                continue
            funcs = []
            for fname, d in fdesc.items():
                if fname in ("components", "Response"):     # doc noise
                    continue
                funcs.append({
                    "name": fname,
                    "description": str(d).split("\n")[0][:300],
                    "parameters": {"type": "dict",
                                   "properties": _toolalpaca_params(str(d)),
                                   "required": []}})
            if not funcs:
                continue
            name = tool.get("Name", "api")
            for i, instr in enumerate(tool.get("Instructions", []) or []):
                if not instr or len(instr) < 8:
                    continue
                tid = f"toolalpaca_{split}_{name}_{i}"
                tasks.append(Task(id=tid, query=instr, functions=funcs,
                                  origin_id=tid))
    return tasks


@dataclass
class Task:
    id: str
    query: str                      # natural-language user request
    functions: list[dict[str, Any]]  # tool schemas offered for this query
    origin_id: str                  # BFCL id this was derived from


def _load_jsonl(path: Path) -> list[dict]:
    return [json.loads(l) for l in path.read_text().splitlines() if l.strip()]


def load_bfcl(categories: tuple[str, ...] = ("simple_python", "multiple",
                                             "parallel")) -> list[Task]:
    fmap = {
        "simple_python": "BFCL_v4_simple_python.json",
        "multiple": "BFCL_v4_multiple.json",
        "parallel": "BFCL_v4_parallel.json",
    }
    tasks: list[Task] = []
    for cat in categories:
        rows = _load_jsonl(DATA_DIR / fmap[cat])
        for r in rows:
            # question is [[{role, content}, ...]] -- take first user turn.
            turns = r["question"][0]
            user_msg = next((m["content"] for m in turns if m["role"] == "user"),
                            turns[0]["content"])
            tasks.append(Task(id=r["id"], query=user_msg,
                              functions=r["function"], origin_id=r["id"]))
    return tasks


_SEAL_TYPE_MAP = {"str": "string", "int": "integer", "float": "number",
                  "bool": "boolean", "list": "array", "dict": "object"}


def _seal_schema(tool: dict) -> dict[str, Any]:
    """Convert one Seal-Tools registry entry to a BFCL-style function schema.

    Keep only {type, description} per property: Seal-Tools specs carry
    non-standard keywords (e.g. stringified enums) that strict JSON-schema
    validators in serving stacks reject.
    """
    props = {}
    for pname, spec in (tool.get("parameters") or {}).items():
        typ = _SEAL_TYPE_MAP.get(spec.get("type", "str"),
                                 spec.get("type", "string"))
        if typ not in ("string", "integer", "number", "boolean",
                       "array", "object"):
            typ = "string"
        clean: dict[str, Any] = {"type": typ,
                                 "description": str(spec.get("description",
                                                             ""))}
        if typ == "array":
            clean["items"] = {"type": "string"}
        props[pname] = clean
    return {
        "name": tool["api_name"],
        "description": tool.get("api_description", ""),
        "parameters": {"type": "dict", "properties": props,
                       "required": tool.get("required", [])},
    }


def _apibank_schema(spec: dict) -> dict[str, Any]:
    props = {}
    for pname, p in (spec.get("input_parameters") or {}).items():
        typ = _SEAL_TYPE_MAP.get(str(p.get("type", "str")).lower(),
                                 str(p.get("type", "string")).lower())
        if typ not in ("string", "integer", "number", "boolean",
                       "array", "object"):
            typ = "string"
        prop: dict[str, Any] = {"type": typ,
                                "description": str(p.get("description", ""))[:200]}
        if typ == "array":
            prop["items"] = {"type": "string"}
        props[pname] = prop
    return {"name": spec["name"], "description": str(spec.get("description", ""))[:300],
            "parameters": {"type": "dict", "properties": props,
                           "required": list(props.keys())}}


def load_apibank(levels: tuple[str, ...] = ("level-1", "level-2")) -> list[Task]:
    """API-Bank (Li et al., 2023, arXiv:2304.08244) as Tasks. Each item's
    `input` dialogue is the query; the API specs embedded in `instruction`
    ('API descriptions: {json}\\n{json}...') become the tool registry.
    Downloaded from HF liminghao1630/API-Bank test-data."""
    tasks: list[Task] = []
    seen = set()
    for lvl in levels:
        f = APIBANK_DIR / f"{lvl}-api.json"
        if not f.exists():
            continue
        for r in json.loads(f.read_text()):
            instr = r.get("instruction", "")
            query = (r.get("input") or "").strip()
            if not query or len(query) < 8:
                continue
            # parse the embedded API-description JSON objects
            funcs = []
            body = instr.split("API descriptions:", 1)
            if len(body) == 2:
                for line in body[1].splitlines():
                    line = line.strip()
                    if line.startswith('{') and '"name"' in line:
                        try:
                            spec = json.loads(line)
                            funcs.append(_apibank_schema(spec))
                        except Exception:
                            pass
            if not funcs:
                continue
            tid = f"apibank_{lvl}_{r.get('file','')}_{r.get('id')}"
            if tid in seen:
                continue
            seen.add(tid)
            tasks.append(Task(id=tid, query=query, functions=funcs,
                              origin_id=tid))
    return tasks


def load_sealtools(n_distractors: int = 3) -> list[Task]:
    """Load the Seal-Tools in-domain test split (Wu et al., 2024) as Tasks.

    Second benchmark for the acceptance experiment: same Task interface as
    load_bfcl(), so the simulated-user construction and the 3-arm replay are
    IDENTICAL to the BFCL runs. Each example keeps its first gold API's schema
    and adds ``n_distractors`` deterministic distractor schemas (stable hash of
    the example id), shuffled deterministically so the gold schema's position
    carries no signal. Only ~49% of Seal-Tools queries contain perturbable
    numerals (vs. most BFCL queries), so the per-user recurrence statistics are
    natively different from BFCL -- more exact repeats, less argument drift.
    """
    import hashlib
    registry = {t["api_name"]: t
                for t in _load_jsonl(SEALTOOLS_DIR / "tool.jsonl")}
    names = sorted(registry)
    tasks: list[Task] = []
    for r in _load_jsonl(SEALTOOLS_DIR / "test_in_domain.jsonl"):
        gold = r["calling"][0]["api"]
        if gold not in registry:      # all resolve in practice; guard anyway
            continue
        h = int(hashlib.sha256(r["id"].encode()).hexdigest(), 16)
        rng = random.Random(h)
        distractors = [n for n in rng.sample(names, n_distractors + 1)
                       if n != gold][:n_distractors]
        funcs = [_seal_schema(registry[n]) for n in [gold] + distractors]
        rng.shuffle(funcs)
        tasks.append(Task(id=r["id"], query=r["query"], functions=funcs,
                          origin_id=r["id"]))
    return tasks


def load_tau2(pool_size: int = 900, seed: int = 20260716) -> list[Task]:
    """Load tau2-bench frozen-trajectory decision points (Barres et al., 2025).

    Third benchmark: multi-turn conversational decision points extracted from
    tau2-bench's shipped reference trajectories (see tau2_extract.py for the
    frozen-conversation design that keeps the 3-arm comparison controlled).
    Each Task's query is a rendered transcript prefix; functions are the full
    tool registry of the task's domain (14 airline / 16 retail tools), so the
    model must also pick the right tool, not just fill arguments. A
    deterministic sample of ``pool_size`` keeps the pool comparable to the
    BFCL (800) / Seal-Tools (700) pools.
    """
    tools = {dom: json.loads((TAU2_DIR / f"tools_{dom}.json").read_text())
             for dom in ("airline", "retail")}
    rows = _load_jsonl(TAU2_DIR / "decision_points.jsonl")
    rng = random.Random(seed)
    rng.shuffle(rows)
    rows = rows[:pool_size]
    return [Task(id=r["id"], query=r["query"],
                 functions=tools[r["domain"]], origin_id=r["id"])
            for r in rows]


_NUM_RE = re.compile(r"(?<![\w.])(-?\d+(?:\.\d+)?)(?![\w.])")


def perturb_numeric(query: str, rng: random.Random) -> str:
    """Replace standalone numbers in a query with new values of similar scale.

    Returns the query unchanged if it contains no substitutable numbers (in
    which case the task recurs as an exact repeat, which is also realistic).
    """
    matches = list(_NUM_RE.finditer(query))
    if not matches:
        return query
    out, last = [], 0
    for m in matches:
        out.append(query[last:m.start()])
        tok = m.group(1)
        if "." in tok:
            base = float(tok)
            lo, hi = max(0.1, base * 0.4), base * 1.9 + 1.0
            out.append(f"{rng.uniform(lo, hi):.1f}")
        else:
            base = int(tok)
            if abs(base) <= 1:            # keep tiny counts (0/1/2) stable
                out.append(tok)
            else:
                lo, hi = max(2, int(abs(base) * 0.4)), int(abs(base) * 1.9) + 2
                val = rng.randint(lo, hi)
                out.append(str(-val if base < 0 else val))
        last = m.end()
    out.append(query[last:])
    return "".join(out)