from __future__ import annotations

import json
import re
from typing import Any

from onefirewall_ai._onefirewall_snapshot import FUNCTION_CALL_XML_TAGS

FUNCTION_CALLS_TAG = FUNCTION_CALL_XML_TAGS["FUNCTION_CALLS_TAG"]
INVOKE_TAG = FUNCTION_CALL_XML_TAGS["INVOKE_TAG"]
PARAMETER_TAG = FUNCTION_CALL_XML_TAGS["PARAMETER_TAG"]
VALUE_TAG = FUNCTION_CALL_XML_TAGS["VALUE_TAG"]

_FUNCTION_CALL_OPEN = f"<{FUNCTION_CALLS_TAG}>"
_FUNCTION_CALL_CLOSE = f"</{FUNCTION_CALLS_TAG}>"
_INVOKE_RE = re.compile(
    rf'<{INVOKE_TAG}\s+name=["\'](.*?)["\']>(.*?)</{INVOKE_TAG}>',
    re.DOTALL | re.IGNORECASE,
)
_PARAMETER_RE = re.compile(
    rf'<{PARAMETER_TAG}\s+name=["\'](.*?)["\']>(.*?)</{PARAMETER_TAG}>',
    re.DOTALL | re.IGNORECASE,
)
_VALUE_WRAPPER_RE = re.compile(rf'</?{VALUE_TAG}\b[^>]*>', re.IGNORECASE)

_TOOL_TAG_PREFIXES = tuple(
    prefix.lower()
    for prefix in (
        f"<{FUNCTION_CALLS_TAG}",
        f"</{FUNCTION_CALLS_TAG}",
        f"<{INVOKE_TAG}",
        f"</{INVOKE_TAG}",
        f"<{PARAMETER_TAG}",
        f"</{PARAMETER_TAG}",
        f"<{VALUE_TAG}",
        f"</{VALUE_TAG}",
    )
)


def _strip_value_wrapper(value: str) -> str:
    return _VALUE_WRAPPER_RE.sub("", value).strip()


def _strip_trailing_partial_tool_tag(text: str) -> str:
    lowered = text.lower()
    cut_index: int | None = None

    for prefix in _TOOL_TAG_PREFIXES:
        index = lowered.rfind(prefix)
        if index == -1:
            continue
        if ">" in text[index:]:
            continue
        cut_index = index if cut_index is None else max(cut_index, index)

    if cut_index is None:
        return text
    return text[:cut_index]


def strip_onefirewall_xml_for_display(text: str) -> str:
    cleaned_text, _ = extract_onefirewall_xml_tool_calls(text)
    return _strip_trailing_partial_tool_tag(cleaned_text)


def extract_onefirewall_xml_tool_calls(text: str) -> tuple[str, list[dict[str, Any]]]:
    if not text:
        return text, []

    lowered = text.lower()
    open_token = _FUNCTION_CALL_OPEN.lower()
    close_token = _FUNCTION_CALL_CLOSE.lower()
    cursor = 0
    cleaned_parts: list[str] = []
    tool_calls: list[dict[str, Any]] = []

    while True:
        start = lowered.find(open_token, cursor)
        if start == -1:
            break

        cleaned_parts.append(text[cursor:start])
        block_start = start + len(_FUNCTION_CALL_OPEN)
        close = lowered.find(close_token, block_start)

        if close == -1:
            block = text[block_start:]
            cursor = len(text)
        else:
            block = text[block_start:close]
            cursor = close + len(_FUNCTION_CALL_CLOSE)

        for match in _INVOKE_RE.finditer(block):
            tool_name = match.group(1).strip()
            invoke_body = match.group(2)
            args: dict[str, str] = {}

            for parameter_match in _PARAMETER_RE.finditer(invoke_body):
                param_name = parameter_match.group(1).strip()
                param_value = _strip_value_wrapper(parameter_match.group(2))
                args[param_name] = param_value

            tool_calls.append(
                {
                    "id": f"call_xml_{len(tool_calls)}",
                    "type": "function",
                    "function": {
                        "name": tool_name,
                        "arguments": json.dumps(args),
                    },
                }
            )

        if close == -1:
            break

    cleaned_parts.append(text[cursor:])
    cleaned_text = "".join(cleaned_parts).strip()
    return cleaned_text, tool_calls
