"""Thin client over the sglang OpenAI-compatible endpoint for gpt-oss-120b.

Given a query and the offered tool schemas, we ask the genuinely-served model
to emit a tool call and return the (name, arguments) it produced. This is the
*target* generation used by the acceptance metric. Generation is greedy
(temperature 0) so the target is deterministic -- the correct reference for a
speculative decoder that verifies greedily.
"""
from __future__ import annotations

import json
import re
import time
from typing import Any

import os

import requests

# The sglang server runs on the GPU node; from the GPU node itself localhost
# works, but from the head/launch node it must be addressed by host. Honor an
# env override so the harness runs from either place.
DEFAULT_URL = os.environ.get("TOOL_SERVER_URL", "http://localhost:30000/v1")

# gpt-oss harmony tool-call markup, e.g.:
#   ...to=functions.triangle_properties.get <|constrain|>json<|message|>{...}<|call|>
_HARMONY_NAME = re.compile(r"to=functions\.([A-Za-z0-9_.\-]+)")
_HARMONY_ARGS = re.compile(r"<\|message\|>(.*?)<\|call\|>", re.DOTALL)


def parse_harmony_content(content: str) -> dict[str, Any] | None:
    """Fallback: extract the first tool call from raw harmony content.

    sglang's HarmonyParser occasionally leaves the call in `content` instead of
    populating `tool_calls`; this recovers the genuine call the model emitted.
    """
    if not content or "to=functions." not in content:
        return None
    nm = _HARMONY_NAME.search(content)
    if not nm:
        return None
    am = _HARMONY_ARGS.search(content, nm.end())
    raw = am.group(1).strip() if am else "{}"
    try:
        args = json.loads(raw) if raw else {}
    except json.JSONDecodeError:
        args = {"__raw__": raw}
    return {"name": nm.group(1), "arguments": args}


# Nemotron-H / Llama-style XML tool-call markup, emitted in `content` when the
# server-side parser (hermes) does not recognize it, e.g.:
#   <tool_call>\n<function=get_weather>\n<parameter=city>\nParis\n</parameter>...
_XML_FUNC = re.compile(r"<function=([A-Za-z0-9_.\-]+)\s*>")
_XML_PARAM = re.compile(r"<parameter=([A-Za-z0-9_.\-]+)\s*>(.*?)</parameter>",
                        re.DOTALL)


def parse_xml_content(content: str) -> dict[str, Any] | None:
    """Fallback: extract an XML-style ``<function=..><parameter=..>`` tool call.

    Used for models (e.g. Nemotron-3-Super) whose native tool-call format the
    served parser leaves in ``content``. Values are kept as the model's literal
    strings; since draft and target pass through the same parser, only their
    mutual agreement matters for the acceptance metric.
    """
    if not content or "<function=" not in content:
        return None
    fm = _XML_FUNC.search(content)
    if not fm:
        return None
    # scope parameters to this function block if a closing tag exists
    end = content.find("</function>", fm.end())
    block = content[fm.end():end if end != -1 else None]
    args: dict[str, Any] = {}
    for pm in _XML_PARAM.finditer(block):
        val = pm.group(2).strip()
        # coerce obvious scalars so canonicalization matches JSON tool_calls
        low = val.lower()
        if low in ("true", "false"):
            args[pm.group(1)] = (low == "true")
        else:
            try:
                args[pm.group(1)] = int(val)
            except ValueError:
                try:
                    args[pm.group(1)] = float(val)
                except ValueError:
                    args[pm.group(1)] = val
    return {"name": fm.group(1), "arguments": args}


# BFCL uses Python-style type names; JSON Schema needs these mappings.
_TYPE_MAP = {"dict": "object", "float": "number", "integer": "integer",
             "tuple": "array", "list": "array", "string": "string",
             "boolean": "boolean", "bool": "boolean", "int": "integer",
             "number": "number", "array": "array", "object": "object"}


def _sanitize_schema(node: Any) -> Any:
    """Recursively convert BFCL Python types into valid JSON Schema."""
    if isinstance(node, dict):
        out = {}
        for k, v in node.items():
            if k == "type" and isinstance(v, str):
                if v == "any":
                    continue                       # unconstrained -> omit type
                out[k] = _TYPE_MAP.get(v, v)
            else:
                out[k] = _sanitize_schema(v)
        # a "tuple"/"array" with no item schema still needs items for strict
        # validators; leave as-is otherwise.
        return out
    if isinstance(node, list):
        return [_sanitize_schema(x) for x in node]
    return node


def to_openai_tools(functions: list[dict[str, Any]]) -> list[dict]:
    """Convert BFCL function schemas to OpenAI tool schema."""
    tools = []
    for f in functions:
        params = f.get("parameters", {}) or {"type": "object", "properties": {}}
        params = _sanitize_schema(dict(params))
        if params.get("type") in (None, "dict"):
            params["type"] = "object"
        tools.append({
            "type": "function",
            "function": {
                "name": f["name"],
                "description": f.get("description", ""),
                "parameters": params,
            },
        })
    return tools


class ToolClient:
    def __init__(self, url: str = DEFAULT_URL, model: str = "gpt-oss-120b",
                 timeout: float = 120.0):
        self.url = url.rstrip("/")
        self.model = model
        self.timeout = timeout

    def ping(self) -> bool:
        try:
            r = requests.get(f"{self.url}/models", timeout=5)
            return r.status_code == 200
        except Exception:
            return False

    def generate_call(self, query: str, functions: list[dict[str, Any]],
                      retries: int = 3) -> dict[str, Any] | None:
        """Return {'name':..., 'arguments':{...}} for the model's tool call.

        Returns None if the model declined to call a tool or on hard failure.
        """
        tools = to_openai_tools(functions)
        payload = {
            "model": self.model,
            "messages": [
                {"role": "system", "content":
                 "You are a function-calling agent. Call exactly one of the "
                 "provided tools to satisfy the user's request."},
                {"role": "user", "content": query},
            ],
            "tools": tools,
            # gpt-oss harmony parser rejects tool_choice="required"
            # (structure_info conflict); it natively emits tool calls with auto.
            "tool_choice": "auto",
            "temperature": 0.0,
            "max_tokens": 512,
        }
        last_err = None
        for attempt in range(retries):
            try:
                r = requests.post(f"{self.url}/chat/completions", json=payload,
                                  timeout=self.timeout)
                if r.status_code != 200:
                    last_err = f"http {r.status_code}: {r.text[:200]}"
                    time.sleep(1.5 * (attempt + 1))
                    continue
                msg = r.json()["choices"][0]["message"]
                tcs = msg.get("tool_calls") or []
                if not tcs:
                    # parser left the call in content -> recover it ourselves.
                    # Try gpt-oss harmony markup first, then Nemotron/Llama XML.
                    content = msg.get("content") or ""
                    return (parse_harmony_content(content)
                            or parse_xml_content(content))
                fn = tcs[0]["function"]
                args = fn.get("arguments", "{}")
                if isinstance(args, str):
                    try:
                        args = json.loads(args) if args.strip() else {}
                    except json.JSONDecodeError:
                        args = {"__raw__": args}
                return {"name": fn["name"], "arguments": args}
            except Exception as e:  # noqa: BLE001
                last_err = str(e)
                time.sleep(1.5 * (attempt + 1))
        raise RuntimeError(f"generate_call failed after {retries}: {last_err}")