import asyncio
import json
import re
from typing import Any, Dict, List, Optional, Protocol


class ToolResult:
    def __init__(self, success: bool, output: str, error: Optional[str] = None):
        self.success = success
        self.output = output
        self.error = error


class Tool(Protocol):
    name: str
    description: str
    parameters: Dict[str, Dict[str, Any]]

    async def execute(self, args: Dict[str, Any]) -> ToolResult:
        ...


class ParsedToolCall:
    __slots__ = ("tool_name", "args")

    def __init__(self, tool_name: str, args: Dict[str, Any]):
        self.tool_name = tool_name
        self.args = args


class ToolRegistry:
    def __init__(self):
        self.tools: Dict[str, Tool] = {}

    def register(self, tool: Tool) -> None:
        self.tools[tool.name] = tool

    def get(self, name: str) -> Optional[Tool]:
        return self.tools.get(name)

    def list(self) -> List[Tool]:
        return list(self.tools.values())

    def describe_for_llm(self) -> str:
        tools = self.list()
        if not tools:
            return ""

        desc = (
            "## Available Tools\n\n"
            "You can call tools by responding with a JSON block. For a **single** tool call:\n"
            '```json\n{"tool": "tool_name", "args": {"param": "value"}}\n```\n\n'
            "For **multiple independent** tool calls that can run in parallel, use an array:\n"
            "```json\n[\n"
            '  {"tool": "read_file", "args": {"path": "a.txt"}},\n'
            '  {"tool": "read_file", "args": {"path": "b.txt"}}\n'
            "]\n```\n\n"
            "Use the array form when the calls are independent (no call depends on another's result).\n\n"
        )

        for tool in tools:
            desc += f"### {tool.name}\n{tool.description}\n"
            params = tool.parameters
            if params:
                desc += "Parameters:\n"
                for name, info in params.items():
                    req = " (required)" if info.get("required") else ""
                    desc += f"- `{name}` ({info.get('type')}{req}): {info.get('description')}\n"
            desc += "\n"

        return desc

    def parse_tool_call(self, response: str) -> Optional[Dict[str, Any]]:
        """Parse a single tool call (legacy compat — returns dict with toolName/args keys)."""
        calls = self.parse_tool_calls(response)
        if not calls:
            return None
        c = calls[0]
        return {"toolName": c.tool_name, "args": c.args}

    def parse_tool_calls(self, response: str) -> List[ParsedToolCall]:
        """Parse one or more tool calls from an LLM response.

        Supports both single-object and array-of-objects JSON formats.
        """
        def try_parse(text: str) -> List[ParsedToolCall]:
            try:
                parsed = json.loads(text)
            except json.JSONDecodeError:
                return []

            if isinstance(parsed, list):
                return [
                    ParsedToolCall(tool_name=item["tool"], args=item.get("args") or {})
                    for item in parsed
                    if isinstance(item, dict) and isinstance(item.get("tool"), str)
                ]
            if isinstance(parsed, dict) and isinstance(parsed.get("tool"), str):
                return [ParsedToolCall(tool_name=parsed["tool"], args=parsed.get("args") or {})]
            return []

        # Try ALL fenced code blocks (not just the first), because the LLM
        # may include non-JSON blocks (e.g. ```python) before the tool call block.
        for json_match in re.finditer(r"```(?:json)?\s*\n?([\s\S]*?)\n?\s*```", response):
            calls = try_parse(json_match.group(1))
            if calls:
                return calls

        # Try direct JSON
        calls = try_parse(response.strip())
        if calls:
            return calls

        return []

    async def execute_tool(self, tool_name: str, args: Dict[str, Any]) -> ToolResult:
        tool = self.get(tool_name)
        if not tool:
            return ToolResult(success=False, output="", error=f"Unknown tool: {tool_name}")
        try:
            return await tool.execute(args)
        except Exception as err:
            return ToolResult(success=False, output="", error=f"Tool error: {str(err)}")

    async def execute_tools_parallel(self, calls: List[ParsedToolCall]) -> List[ToolResult]:
        """Execute multiple tool calls concurrently using asyncio.gather.

        Returns results in the same order as the input calls.
        """
        if not calls:
            return []
        if len(calls) == 1:
            return [await self.execute_tool(calls[0].tool_name, calls[0].args)]

        async def _safe_exec(call: ParsedToolCall) -> ToolResult:
            try:
                return await self.execute_tool(call.tool_name, call.args)
            except Exception as err:
                return ToolResult(success=False, output="", error=f"Tool error: {str(err)}")

        return list(await asyncio.gather(*[_safe_exec(c) for c in calls]))
