"""
Universal tool abstraction for Daita agents.

Tools are LLM-callable functions that can come from:
- Plugins (database queries, S3 operations, API calls)
- MCP servers (external tools via Model Context Protocol)
- Custom functions (user-defined Python functions)

This abstraction is provider-agnostic and supports both prompt-based
and native function calling (OpenAI, Anthropic, etc).
"""

from dataclasses import dataclass
from typing import Callable, Dict, Any, List, Optional, Awaitable, Union
import asyncio
import inspect
import logging

logger = logging.getLogger(__name__)


@dataclass
class AgentTool:
    """
    Universal tool definition for agent-LLM integration.

    Represents any callable function that an agent can use, regardless of source.
    Designed to be compatible with industry standards (LangChain, AutoGen, etc).

    Example:
        ```python
        tool = AgentTool(
            name="search_database",
            description="Search for records in the database",
            parameters={
                "query": {
                    "type": "string",
                    "description": "SQL query to execute",
                    "required": True
                }
            },
            handler=async_search_function
        )
        ```
    """

    # Core fields (required)
    name: str
    description: str
    parameters: Dict[str, Any]  # JSON Schema format
    handler: Callable[[Dict[str, Any]], Awaitable[Any]]  # async function

    # Optional metadata
    category: Optional[str] = None  # "database", "storage", "api", etc
    source: str = "custom"  # "plugin", "mcp", "custom"
    plugin_name: Optional[str] = None  # Which plugin provides this tool

    # Safety features
    timeout_seconds: Optional[int] = None  # Execution timeout

    def to_openai_function(self) -> Dict[str, Any]:
        """
        Convert to OpenAI function calling format.

        Returns:
            OpenAI function definition dict

        Reference:
            https://platform.openai.com/docs/guides/function-calling
        """
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": {
                    "type": "object",
                    "properties": self.parameters,
                    "required": [
                        k for k, v in self.parameters.items()
                        if v.get("required", False)
                    ]
                }
            }
        }

    def to_anthropic_tool(self) -> Dict[str, Any]:
        """
        Convert to Anthropic tool format.

        Returns:
            Anthropic tool definition dict

        Reference:
            https://docs.anthropic.com/claude/docs/tool-use
        """
        return {
            "name": self.name,
            "description": self.description,
            "input_schema": {
                "type": "object",
                "properties": self.parameters,
                "required": [
                    k for k, v in self.parameters.items()
                    if v.get("required", False)
                ]
            }
        }

    def to_llm_function(self) -> Dict[str, Any]:
        """
        Generic LLM function format (works for most providers).

        For provider-specific formats, use to_openai_function() or to_anthropic_tool().
        This format is used for prompt-based tool calling.
        """
        return {
            "name": self.name,
            "description": self.description,
            "parameters": {
                "type": "object",
                "properties": self.parameters,
                "required": [
                    k for k, v in self.parameters.items()
                    if v.get("required", False)
                ]
            }
        }

    def to_prompt_description(self) -> str:
        """
        Generate human-readable tool description for prompt injection.

        Used in prompt-based tool calling (non-native function calling).

        Returns:
            Formatted tool description string
        """
        params_desc = []
        for param_name, param_info in self.parameters.items():
            required = " (required)" if param_info.get("required") else " (optional)"
            param_type = param_info.get("type", "any")
            param_desc = param_info.get("description", "")
            params_desc.append(
                f"  - {param_name} ({param_type}){required}: {param_desc}"
            )

        params_str = "\n".join(params_desc) if params_desc else "  (no parameters)"

        return f"{self.name}: {self.description}\nParameters:\n{params_str}"

    async def execute(self, arguments: Dict[str, Any]) -> Any:
        """
        Execute the tool with given arguments.

        Args:
            arguments: Tool arguments matching the parameter schema

        Returns:
            Tool execution result

        Raises:
            RuntimeError: If tool execution fails or times out
        """
        if not callable(self.handler):
            raise RuntimeError(f"Tool '{self.name}' has non-callable handler")

        # Execute with timeout if specified
        if self.timeout_seconds:
            try:
                result = await asyncio.wait_for(
                    self.handler(arguments),
                    timeout=self.timeout_seconds
                )
                return result
            except asyncio.TimeoutError:
                raise RuntimeError(
                    f"Tool '{self.name}' execution timed out after {self.timeout_seconds}s"
                )
        else:
            # No timeout
            return await self.handler(arguments)

    @classmethod
    def from_function(
        cls,
        func: Callable,
        name: Optional[str] = None,
        description: Optional[str] = None,
        parameters: Optional[Dict[str, Any]] = None,
        **kwargs
    ) -> 'AgentTool':
        """
        Create AgentTool from a Python function.

        Automatically extracts name, docstring, and parameter info if not provided.

        Args:
            func: Python function (sync or async)
            name: Tool name (defaults to function name)
            description: Tool description (defaults to docstring)
            parameters: Parameter schema (auto-detected from type hints if not provided)
            **kwargs: Additional AgentTool fields (category, timeout_seconds, etc)

        Returns:
            AgentTool instance

        Example:
            ```python
            async def search_docs(query: str, limit: int = 10):
                '''Search documentation with a query'''
                return search_results

            # Auto-detection from function signature
            tool = AgentTool.from_function(search_docs)

            # Or with explicit parameters
            tool = AgentTool.from_function(
                search_docs,
                parameters={
                    "query": {
                        "type": "string",
                        "description": "Search query",
                        "required": True
                    },
                    "limit": {
                        "type": "integer",
                        "description": "Max results",
                        "required": False
                    }
                }
            )
            ```
        """
        # Use function name if not provided
        tool_name = name or func.__name__

        # Use docstring if no description
        tool_description = description or (func.__doc__ or f"Execute {tool_name}").strip()

        # Auto-detect parameters from type hints if not provided
        if parameters is None:
            parameters = {}
            sig = inspect.signature(func)

            for param_name, param in sig.parameters.items():
                # Skip self/cls parameters
                if param_name in ('self', 'cls'):
                    continue

                # Map Python types to JSON Schema types
                param_type = "string"  # Default type
                if param.annotation != inspect.Parameter.empty:
                    if param.annotation in (int, 'int'):
                        param_type = "integer"
                    elif param.annotation in (float, 'float'):
                        param_type = "number"
                    elif param.annotation in (bool, 'bool'):
                        param_type = "boolean"
                    elif param.annotation in (list, List):
                        param_type = "array"
                    elif param.annotation in (dict, Dict):
                        param_type = "object"

                # Determine if required (no default value)
                is_required = param.default == inspect.Parameter.empty

                parameters[param_name] = {
                    "type": param_type,
                    "description": f"Parameter: {param_name}",
                    "required": is_required
                }

        # Wrap sync functions to make them async
        if asyncio.iscoroutinefunction(func):
            handler = func
        else:
            async def async_wrapper(args: Dict[str, Any]) -> Any:
                return func(**args)
            handler = async_wrapper

        return cls(
            name=tool_name,
            description=tool_description,
            parameters=parameters,
            handler=handler,
            source="custom",
            **kwargs
        )

    @classmethod
    def from_mcp_tool(cls, mcp_tool, mcp_registry) -> 'AgentTool':
        """
        Create AgentTool from an MCP tool.

        Args:
            mcp_tool: MCPTool instance from MCP plugin
            mcp_registry: MCPToolRegistry for routing calls

        Returns:
            AgentTool instance that wraps the MCP tool
        """
        # Create handler that routes to MCP registry
        async def mcp_handler(arguments: Dict[str, Any]) -> Any:
            return await mcp_registry.call_tool(mcp_tool.name, arguments)

        return cls(
            name=mcp_tool.name,
            description=mcp_tool.description,
            parameters=mcp_tool.input_schema.get("properties", {}),
            handler=mcp_handler,
            source="mcp",
            category="mcp"
        )


class ToolRegistry:
    """
    Registry for managing tools from multiple sources.

    Used internally by agents to aggregate tools from plugins, MCP servers,
    and custom functions.
    """

    def __init__(self):
        """Initialize empty tool registry"""
        self.tools: List[AgentTool] = []
        self._tool_map: Dict[str, AgentTool] = {}

    def register(self, tool: AgentTool) -> None:
        """
        Register a tool.

        Args:
            tool: AgentTool to register
        """
        if tool.name in self._tool_map:
            logger.warning(
                f"Tool '{tool.name}' already registered. "
                f"Overwriting (old source: {self._tool_map[tool.name].source}, "
                f"new source: {tool.source})"
            )

        self.tools.append(tool)
        self._tool_map[tool.name] = tool

        logger.debug(f"Registered tool: {tool.name} (source: {tool.source})")

    def register_many(self, tools: List[AgentTool]) -> None:
        """
        Register multiple tools at once.

        Args:
            tools: List of AgentTool instances
        """
        for tool in tools:
            self.register(tool)

    def get(self, name: str) -> Optional[AgentTool]:
        """
        Get tool by name.

        Args:
            name: Tool name

        Returns:
            AgentTool instance or None if not found
        """
        return self._tool_map.get(name)

    async def execute(self, name: str, arguments: Dict[str, Any]) -> Any:
        """
        Execute a tool by name.

        Args:
            name: Tool name
            arguments: Tool arguments

        Returns:
            Tool execution result

        Raises:
            RuntimeError: If tool not found or execution fails
        """
        tool = self.get(name)
        if not tool:
            available = [t.name for t in self.tools]
            raise RuntimeError(
                f"Tool '{name}' not found. Available tools: {', '.join(available)}"
            )

        return await tool.execute(arguments)

    @property
    def tool_count(self) -> int:
        """Total number of registered tools"""
        return len(self.tools)

    @property
    def tool_names(self) -> List[str]:
        """List of all tool names"""
        return list(self._tool_map.keys())
