from ast import literal_eval
from ..model.entities import ToolCall, ToolFormat
from json import JSONDecodeError, loads
from re import DOTALL, search
from typing import Any, Optional, Tuple
from xml.etree import ElementTree


class ToolCallParser:
    _eos_token: Optional[str]
    _tool_format: Optional[ToolFormat]

    def __init__(
        self,
        tool_format: Optional[ToolFormat] = None,
        eos_token: Optional[str] = None,
    ) -> None:
        self._tool_format = tool_format
        self._eos_token = eos_token

    def __call__(self, text: str) -> Optional[list[ToolCall]]:
        calls = (
            self._parse_json(text)
            if self._tool_format is ToolFormat.JSON
            else self._parse_react(text)
            if self._tool_format is ToolFormat.REACT
            else self._parse_bracket(text)
            if self._tool_format is ToolFormat.BRACKET
            else self._parse_openai_json(text)
            if self._tool_format is ToolFormat.OPENAI
            else None
        )
        if not calls:
            calls = self._parse_tag(text)
        return calls

    def set_eos_token(self, eos_token: str) -> None:
        self._eos_token = eos_token

    def _parse_json(self, text: str) -> Optional[Tuple[str, dict[str, Any]]]:
        try:
            payload = loads(text)
            return payload["tool"], payload.get("arguments", {})
        except Exception:
            return None

    def _parse_react(self, text: str) -> Optional[Tuple[str, dict[str, Any]]]:
        act = search(r"Action:\s*(\w+)", text)
        inp = search(r"Action Input:\s*({.*})", text, DOTALL)
        if act and inp:
            try:
                return act.group(1), loads(inp.group(1))
            except JSONDecodeError:
                pass
        return None

    def _parse_bracket(self, text: str) -> Optional[Tuple[str, dict[str, Any]]]:
        m = search(r"\[(\w+)\]\(([^)]+)\)", text)
        if m:
            return m.group(1), {"input": m.group(2)}
        return None

    def _parse_openai_json(
        self, text: str
    ) -> Optional[Tuple[str, dict[str, Any]]]:
        try:
            payload = loads(text)
            name = payload.get("name")
            args = payload.get("arguments", {})
            if isinstance(name, str) and isinstance(args, dict):
                return name, args
        except JSONDecodeError:
            pass
        return None

    def _parse_tag(self, text: str) -> Optional[Tuple[str, dict[str, Any]]]:
        tool_calls: list[ToolCall] = []

        if self._eos_token:
            text = text.strip().removesuffix(self._eos_token)
        try:
            root = ElementTree.fromstring(f"<root>{text}</root>")
            for element in root.findall(".//tool_call"):
                tool_call = None
                try:
                    json_text = element.text.strip()

                    try:
                        tool_call = loads(json_text)
                    except JSONDecodeError:
                        try:
                            tool_call = literal_eval(json_text)
                        except (SyntaxError, ValueError):
                            continue
                except Exception:
                    pass

                if (
                    tool_call is not None
                    and "name" in tool_call
                    and "arguments" in tool_call
                ):
                    tool_calls.append(
                        ToolCall(
                            name=tool_call["name"],
                            arguments=tool_call["arguments"],
                        )
                    )
        except ElementTree.ParseError:
            pass

        if tool_calls:
            return tool_calls

        m = search(
            r"<tool_call>\s*(\{.*?\})\s*</tool_call>", text, flags=DOTALL
        )
        if m:
            tool_call_payload = m.group(1)
            try:
                tool_call = loads(tool_call_payload)
                if (
                    tool_call
                    and "name" in tool_call
                    and "arguments" in tool_call
                ):
                    tool_calls.append(
                        ToolCall(
                            name=tool_call["name"],
                            arguments=tool_call["arguments"],
                        )
                    )
            except JSONDecodeError:
                pass

        m = search(
            r"<tool_call\s+name=\"([^\"]+)\"\s*>(\{.*?\})</tool_call>",
            text,
            DOTALL,
        )
        if m:
            try:
                tool_calls.append(
                    ToolCall(name=m.group(1), arguments=loads(m.group(2)))
                )
            except JSONDecodeError:
                pass

        m = search(
            r"<tool\s+name=\"([^\"]+)\"\s*>(\{.*?\})</tool>",
            text,
            DOTALL,
        )
        if m:
            try:
                tool_calls.append(
                    ToolCall(name=m.group(1), arguments=loads(m.group(2)))
                )
            except JSONDecodeError:
                pass

        m = search(
            r"<tool_call\s+name=\"([^\"]+)\"\s+arguments='(\{.*?\})'\s*/>",
            text,
        )
        if m:
            try:
                tool_calls.append(
                    ToolCall(name=m.group(1), arguments=loads(m.group(2)))
                )
            except JSONDecodeError:
                pass

        return tool_calls if tool_calls else None
