src/managers/tools/tool_manager.py

"""
Tool Manager for CodeAct Agent.
Unified management system for all types of tools: local functions, decorated tools, and MCP tools.
"""

import re
from collections.abc import Callable
from dataclasses import dataclass
from enum import Enum
from typing import Any

from .mcp_manager import MCPManager

# Import components
from .tool_registry import ToolRegistry


class ToolSource(Enum):
    """Enumeration of tool sources."""

    LOCAL = "local"
    DECORATED = "decorated"
    MCP = "mcp"
    ALL = "all"


@dataclass
class ToolInfo:
    """Comprehensive tool information."""

    name: str
    description: str
    source: ToolSource
    function: Callable | None = None
    schema: dict | None = None
    server: str | None = None  # For MCP tools
    module: str | None = None
    required_parameters: list[dict] = None
    optional_parameters: list[dict] = None

    def __post_init__(self):
        if self.required_parameters is None:
            self.required_parameters = []
        if self.optional_parameters is None:
            self.optional_parameters = []


class ToolManager:
    """
    Unified tool management system for CodeAct Agent.

    Manages all types of tools:
    - Local functions (legacy function registry)
    - Decorated tools (@tool decorator)
    - MCP tools (Model Context Protocol)

    Provides a single, consistent interface for tool operations.
    """

    def __init__(self, console_display=None):
        """
        Initialize the ToolManager.

        Args:
            console_display: Optional console display for MCP status output
        """
        # Core components
        self.tool_registry = ToolRegistry()
        self.mcp_manager = MCPManager(console_display)

        # Unified tool catalog
        self._tool_catalog: dict[str, ToolInfo] = {}

        # Initialize with decorated tools from builtin_tools.py
        self._discover_decorated_tools()

    # ====================
    # CORE TOOL MANAGEMENT
    # ====================

    def add_tool(
        self,
        tool: Callable | dict,
        name: str = None,
        description: str = None,
        source: ToolSource = ToolSource.LOCAL,
    ) -> bool:
        """
        Add a tool to the manager.

        Args:
            tool: Either a callable function or a tool schema dict
            name: Optional custom name (defaults to function.__name__)
            description: Optional description (defaults to function.__doc__)
            source: Source type of the tool

        Returns:
            True if successfully added
        """
        try:
            if callable(tool):
                # Handle callable functions
                tool_name = name or tool.__name__
                tool_desc = description or tool.__doc__ or f"Function {tool.__name__}"

                # Add to tool registry
                success = self.tool_registry.add_function_directly(tool_name, tool, tool_desc)

                if success:
                    # Create ToolInfo and add to catalog
                    tool_info = ToolInfo(
                        name=tool_name,
                        description=tool_desc,
                        source=source,
                        function=tool,
                        schema=self._create_schema_from_function(tool_name, tool, tool_desc),
                    )
                    self._tool_catalog[tool_name] = tool_info

                    return True

            elif isinstance(tool, dict):
                # Handle tool schema dictionaries
                tool_name = tool.get("name") or name
                tool_desc = tool.get("description") or description

                if not tool_name:
                    print("Warning: Tool schema must have a name")
                    return False

                # Register with tool registry
                success = self.tool_registry.register_tool(tool)

                if success:
                    tool_info = ToolInfo(
                        name=tool_name,
                        description=tool_desc,
                        source=source,
                        schema=tool,
                        required_parameters=tool.get("required_parameters", []),
                        optional_parameters=tool.get("optional_parameters", []),
                    )
                    self._tool_catalog[tool_name] = tool_info
                    return True

            return False

        except Exception as e:
            print(f"Error adding tool {name}: {e}")
            return False

    def remove_tool(self, name: str) -> bool:
        """Remove a tool by name from all registries."""
        try:
            success = False

            # Remove from tool registry
            if self.tool_registry.remove_tool_by_name(name):
                success = True

            # Remove from MCP if it's an MCP tool
            if name in self.mcp_manager.mcp_functions:
                if self.mcp_manager.remove_mcp_tool(name):
                    success = True

            # Remove from catalog
            if name in self._tool_catalog:
                del self._tool_catalog[name]
                success = True

            return success

        except Exception as e:
            print(f"Error removing tool {name}: {e}")
            return False

    def get_tool(self, name: str) -> ToolInfo | None:
        """Get comprehensive tool information by name."""
        return self._tool_catalog.get(name)

    def get_tool_function(self, name: str) -> Callable | None:
        """Get the actual function object by name."""
        tool_info = self._tool_catalog.get(name)
        return tool_info.function if tool_info else None

    # ====================
    # TOOL DISCOVERY AND LISTING
    # ====================

    def list_tools(
        self, source: ToolSource = ToolSource.ALL, include_details: bool = False
    ) -> list[dict]:
        """
        List tools with optional filtering by source.

        Args:
            source: Filter by tool source (LOCAL, DECORATED, MCP, ALL)
            include_details: Whether to include detailed information

        Returns:
            List of tool dictionaries
        """
        tools = []

        for _tool_name, tool_info in self._tool_catalog.items():
            if source == ToolSource.ALL or tool_info.source == source:
                if include_details:
                    tools.append(
                        {
                            "name": tool_info.name,
                            "description": tool_info.description,
                            "source": tool_info.source.value,
                            "server": tool_info.server,
                            "module": tool_info.module,
                            "has_function": tool_info.function is not None,
                            "required_params": len(tool_info.required_parameters),
                            "optional_params": len(tool_info.optional_parameters),
                        }
                    )
                else:
                    tools.append(
                        {
                            "name": tool_info.name,
                            "description": tool_info.description,
                            "source": tool_info.source.value,
                        }
                    )

        return sorted(tools, key=lambda x: x["name"])

    def search_tools(self, query: str, search_descriptions: bool = True) -> list[dict]:
        """
        Search tools by name and optionally description.

        Args:
            query: Search query (supports regex)
            search_descriptions: Whether to also search in descriptions

        Returns:
            List of matching tools
        """
        pattern = re.compile(query, re.IGNORECASE)
        matching_tools = []

        for tool_name, tool_info in self._tool_catalog.items():
            match = False

            # Search in name
            if pattern.search(tool_name):
                match = True

            # Search in description if enabled
            elif search_descriptions and pattern.search(tool_info.description or ""):
                match = True

            if match:
                matching_tools.append(
                    {
                        "name": tool_info.name,
                        "description": tool_info.description,
                        "source": tool_info.source.value,
                        "server": tool_info.server,
                    }
                )

        return sorted(matching_tools, key=lambda x: x["name"])

    def get_tools_by_source(self, source: ToolSource) -> dict[str, ToolInfo]:
        """Get all tools from a specific source."""
        return {
            name: tool_info
            for name, tool_info in self._tool_catalog.items()
            if tool_info.source == source
        }

    # ====================
    # MCP INTEGRATION
    # ====================

    def add_mcp_server(self, config_path: str = "./mcp_config.yaml") -> None:
        """Add MCP tools from configuration file."""
        try:
            # Use MCP manager to load tools
            self.mcp_manager.add_mcp(config_path)

            # Update our catalog with MCP tools
            mcp_tools = self.mcp_manager.list_mcp_tools()
            for tool_name, tool_data in mcp_tools.items():
                tool_info = ToolInfo(
                    name=tool_name,
                    description=tool_data.get("description", "MCP tool"),
                    source=ToolSource.MCP,
                    function=tool_data.get("function"),
                    server=tool_data.get("server"),
                    module=tool_data.get("module"),
                    required_parameters=tool_data.get("required_parameters", []),
                    optional_parameters=tool_data.get("optional_parameters", []),
                )
                self._tool_catalog[tool_name] = tool_info

        except Exception as e:
            print(f"Error adding MCP server: {e}")
            # discover_mcp_tools_sync() can succeed (no entry in
            # last_discovery_errors) while a *later* step here — e.g.
            # registering a tool in MCPManager.add_mcp(), or building
            # ToolInfo above — raises and is only ever printed to stdout
            # (invisible on HF Spaces, see 403 on fetch_space_logs). Record
            # it in last_discovery_errors so gradio_ui's health check can
            # surface it, along with how many mcp_functions existed at the
            # point of failure (partial registration vs. none at all).
            import traceback

            self.mcp_manager.last_discovery_errors["_add_mcp_server"] = (
                f"{type(e).__name__}: {e}; "
                f"mcp_functions at failure: {len(self.mcp_manager.mcp_functions)}; "
                f"{traceback.format_exc(limit=4)}"
            )

    def add_mcp_http_server(self, url: str, server_name: str = "decouplerpy") -> None:
        """Register MCP tools from a persistent HTTP MCP server.

        Mirror of add_mcp_server() but for a long-lived `server.py --transport
        http` process: tool calls reuse the running server over HTTP instead of
        spawning a fresh `python server.py` subprocess per call.
        """
        try:
            self.mcp_manager.add_mcp_http(url, server_name=server_name)

            mcp_tools = self.mcp_manager.list_mcp_tools()
            for tool_name, tool_data in mcp_tools.items():
                tool_info = ToolInfo(
                    name=tool_name,
                    description=tool_data.get("description", "MCP tool"),
                    source=ToolSource.MCP,
                    function=tool_data.get("function"),
                    server=tool_data.get("server"),
                    module=tool_data.get("module"),
                    required_parameters=tool_data.get("required_parameters", []),
                    optional_parameters=tool_data.get("optional_parameters", []),
                )
                self._tool_catalog[tool_name] = tool_info
        except Exception as e:
            print(f"Error adding MCP HTTP server: {e}")
            import traceback

            self.mcp_manager.last_discovery_errors["_add_mcp_http_server"] = (
                f"{type(e).__name__}: {e}; "
                f"mcp_functions at failure: {len(self.mcp_manager.mcp_functions)}; "
                f"{traceback.format_exc(limit=4)}"
            )

    def list_mcp_servers(self) -> dict[str, list[str]]:
        """List all MCP servers and their tools."""
        mcp_tools = self.get_tools_by_source(ToolSource.MCP)
        servers = {}

        for tool_name, tool_info in mcp_tools.items():
            server_name = tool_info.server or "unknown"
            if server_name not in servers:
                servers[server_name] = []
            servers[server_name].append(tool_name)

        return servers

    def show_mcp_status(self) -> None:
        """Display detailed MCP status."""
        self.mcp_manager.show_mcp_status()

    def get_mcp_summary(self) -> dict[str, Any]:
        """Get MCP tools summary."""
        return self.mcp_manager.get_mcp_summary()

    # ====================
    # TOOL EXECUTION SUPPORT
    # ====================

    def get_all_functions(self) -> dict[str, Callable]:
        """Get all available functions. Reads from _tool_catalog — the single source of truth."""
        return {
            name: info.function
            for name, info in self._tool_catalog.items()
            if info.function is not None
        }

    def get_tool_schemas(self, openai_format: bool = True) -> list[dict]:
        """
        Get tool schemas for all tools.

        Args:
            openai_format: Whether to format as OpenAI function schemas

        Returns:
            List of tool schemas
        """
        schemas = []

        for _tool_name, tool_info in self._tool_catalog.items():
            if openai_format:
                # Convert to OpenAI function schema format
                schema = {
                    "type": "function",
                    "function": {
                        "name": tool_info.name,
                        "description": tool_info.description,
                        "parameters": {"type": "object", "properties": {}, "required": []},
                    },
                }

                # Add required parameters
                for param in tool_info.required_parameters:
                    param_schema = {
                        "type": param.get("type", "string"),
                        "description": param.get("description", ""),
                    }
                    # Add enum values if present
                    if "enum" in param:
                        param_schema["enum"] = param["enum"]
                    schema["function"]["parameters"]["properties"][param["name"]] = param_schema
                    schema["function"]["parameters"]["required"].append(param["name"])

                # Add optional parameters
                for param in tool_info.optional_parameters:
                    param_schema = {
                        "type": param.get("type", "string"),
                        "description": param.get("description", ""),
                    }
                    # Add enum values if present
                    if "enum" in param:
                        param_schema["enum"] = param["enum"]
                    if "default" in param:
                        param_schema["default"] = param["default"]
                    schema["function"]["parameters"]["properties"][param["name"]] = param_schema

                schemas.append(schema)
            else:
                # Return raw schema
                if tool_info.schema:
                    schemas.append(tool_info.schema)

        return schemas

    # ====================
    # STATISTICS AND REPORTING
    # ====================

    def get_tool_statistics(self) -> dict[str, Any]:
        """Get comprehensive tool statistics."""
        stats = {
            "total_tools": len(self._tool_catalog),
            "by_source": {source.value: 0 for source in ToolSource if source != ToolSource.ALL},
            "with_functions": 0,
            "mcp_servers": len(self.list_mcp_servers()),
            "tool_registry_size": len(self.tool_registry.tools),
        }

        for tool_info in self._tool_catalog.values():
            stats["by_source"][tool_info.source.value] += 1
            if tool_info.function:
                stats["with_functions"] += 1

        return stats

    def validate_tools(self) -> dict[str, list[str]]:
        """Validate all tools and return any issues found."""
        issues = {
            "missing_functions": [],
            "missing_descriptions": [],
            "duplicate_names": [],
            "invalid_schemas": [],
        }

        seen_names = set()
        for tool_name, tool_info in self._tool_catalog.items():
            # Check for duplicates
            if tool_name in seen_names:
                issues["duplicate_names"].append(tool_name)
            seen_names.add(tool_name)

            # Check for missing functions (except MCP tools which may not have direct functions)
            if not tool_info.function and tool_info.source != ToolSource.MCP:
                issues["missing_functions"].append(tool_name)

            # Check for missing descriptions
            if not tool_info.description or tool_info.description.strip() == "":
                issues["missing_descriptions"].append(tool_name)

        return issues

    # ====================
    # PRIVATE METHODS
    # ====================

    def _discover_decorated_tools(self) -> None:
        """Discover and register tools marked with @tool decorator."""
        try:
            from .builtin_tools import get_all_tool_functions

            tool_functions = get_all_tool_functions()
            for func in tool_functions:
                name = getattr(func, "_tool_name", func.__name__)
                description = getattr(
                    func, "_tool_description", func.__doc__ or f"Function {func.__name__}"
                )

                tool_info = ToolInfo(
                    name=name,
                    description=description,
                    source=ToolSource.DECORATED,
                    function=func,
                    schema=self._create_schema_from_function(name, func, description),
                )
                self._tool_catalog[name] = tool_info

                # Also add to tool registry for consistency
                self.tool_registry.add_function_directly(name, func, description)

        except ImportError:
            print("Warning: Could not import builtin_tools module for decorated tool discovery")

    def _create_schema_from_function(self, name: str, function: Callable, description: str) -> dict:
        """Create a tool schema from a function object."""
        return self.tool_registry._create_schema_from_function(name, function, description)