# Copyright (c) 2023-2025, HaiyangLi <quantocean.li at gmail dot com>
# SPDX-License-Identifier: Apache-2.0

import logging
from typing import TYPE_CHECKING, Any

from pydantic import BaseModel

if TYPE_CHECKING:
    from lionagi.service.connections.mcp_wrapper import MCPSecurityConfig

from lionagi.protocols._concepts import Manager
from lionagi.protocols.messages.action_request import ActionRequest
from lionagi.utils import to_list

from .function_calling import FunctionCalling
from .tool import FuncTool, FuncToolRef, Tool, ToolRef

logger = logging.getLogger(__name__)


class ActionManager(Manager):
    """Registers function-based tools and invokes them from ActionRequests."""

    def __init__(self, *args: FuncTool, **kwargs) -> None:
        super().__init__()
        self.registry: dict[str, Tool] = {}

        tools = []
        if args:
            tools.extend(to_list(args, dropna=True, flatten=True))
        if kwargs:
            tools.extend(to_list(kwargs.values(), dropna=True, flatten=True))

        self.register_tools(tools, update=True)

    def __contains__(self, tool: FuncToolRef) -> bool:
        if isinstance(tool, Tool):
            return tool.function in self.registry
        elif isinstance(tool, str):
            return tool in self.registry
        elif callable(tool):
            return tool.__name__ in self.registry
        return False

    def register_tool(self, tool: FuncTool, update: bool = False) -> None:
        if not update and tool in self:
            name = None
            if isinstance(tool, Tool):
                name = tool.function
            elif callable(tool):
                name = tool.__name__
            elif isinstance(tool, dict):
                # For MCP config, extract the tool name (first key)
                name = list(tool.keys())[0] if tool else None
            raise ValueError(f"Tool {name} is already registered.")

        # Convert to Tool object based on type
        if callable(tool):
            tool = Tool(func_callable=tool)
        elif isinstance(tool, dict):
            tool = Tool(mcp_config=tool)
        elif not isinstance(tool, Tool):
            raise TypeError(
                "Must provide a `Tool` object, a callable function, or an MCP config dict."
            )
        self.registry[tool.function] = tool

    def register_tools(self, tools: list[FuncTool] | FuncTool, update: bool = False) -> None:
        tools_list = tools if isinstance(tools, list) else [tools]
        for t in tools_list:
            self.register_tool(t, update=update)

    def match_tool(self, action_request: ActionRequest | BaseModel | dict) -> FunctionCalling:
        if not isinstance(action_request, ActionRequest | BaseModel | dict):
            raise TypeError(f"Unsupported type {type(action_request)}")

        func, args = None, None
        if isinstance(action_request, dict):
            func = action_request["function"]
            args = action_request["arguments"]
        else:
            func = action_request.function
            args = action_request.arguments

        tool = self.registry.get(func, None)
        if not isinstance(tool, Tool):
            raise ValueError(f"Function {func} is not registered.")

        return FunctionCalling(func_tool=tool, arguments=args)

    async def invoke(
        self,
        func_call: BaseModel | ActionRequest,
    ) -> FunctionCalling:
        function_calling = self.match_tool(func_call)
        await function_calling.invoke()
        return function_calling

    @property
    def schema_list(self) -> list[dict[str, Any]]:
        return [tool.tool_schema for tool in self.registry.values()]

    def get_tool_schema(
        self,
        tools: ToolRef = False,
        auto_register: bool = True,
        update: bool = False,
    ) -> dict:
        if isinstance(tools, list | tuple) and len(tools) == 1:
            tools = tools[0]
        if isinstance(tools, bool):
            if tools is True:
                return {"tools": self.schema_list}
            return []
        else:
            schemas = self._get_tool_schema(tools, auto_register=auto_register, update=update)
            return {"tools": schemas}

    def _get_tool_schema(
        self,
        tool: Any,
        auto_register: bool = True,
        update: bool = False,
    ) -> list[dict[str, Any]] | dict[str, Any]:
        if isinstance(tool, dict):
            return tool
        if callable(tool):
            name = tool.__name__
            if name not in self.registry:
                if auto_register:
                    self.register_tool(tool, update=update)
                else:
                    raise ValueError(f"Tool {name} is not registered.")
            return self.registry[name].tool_schema

        elif isinstance(tool, Tool) or isinstance(tool, str):
            name = tool.function if isinstance(tool, Tool) else tool
            if name in self.registry:
                return self.registry[name].tool_schema
            raise ValueError(f"Tool {name} is not registered.")
        elif isinstance(tool, list):
            return [self._get_tool_schema(t, auto_register=auto_register) for t in tool]
        raise TypeError(f"Unsupported type {type(tool)}")

    async def register_mcp_server(
        self,
        server_config: dict[str, Any],
        tool_names: list[str] | None = None,
        request_options: dict[str, type] | None = None,
        update: bool = False,
        security: "MCPSecurityConfig | None" = None,
    ) -> list[str]:
        registered_tools = []

        if security is not None:
            from lionagi.service.connections.mcp_wrapper import MCPConnectionPool

            MCPConnectionPool.remember_security(server_config, security)

        server_name = None
        if isinstance(server_config, dict) and "server" in server_config:
            server_name = server_config["server"]

        if request_options:
            for k in list(request_options.keys()):
                if not k.startswith(f"{server_name}_"):
                    request_options[f"{server_name}_{k}"] = request_options.pop(k)

        if tool_names:
            for tool_name in tool_names:
                config_with_metadata = dict(server_config)
                config_with_metadata["_original_tool_name"] = tool_name

                mcp_config = {tool_name: config_with_metadata}

                tool_request_options = None
                if request_options and tool_name in request_options:
                    tool_request_options = request_options[tool_name]

                tool = Tool(mcp_config=mcp_config, request_options=tool_request_options)
                self.register_tool(tool, update=update)
                registered_tools.append(tool_name)
        else:
            # Auto-discover tools from the server
            from lionagi.service.connections.mcp_wrapper import MCPConnectionPool

            # Get client and discover tools
            client = await MCPConnectionPool.get_client(server_config, security=security)
            tools = await client.list_tools()

            # Register each discovered tool with qualified name
            for tool in tools:
                # Store original tool name in config for MCP calls
                config_with_metadata = dict(server_config)
                config_with_metadata["_original_tool_name"] = tool.name

                mcp_config = {tool.name: config_with_metadata}

                tool_request_options = None
                if request_options and tool.name in request_options:
                    tool_request_options = request_options[tool.name]

                tool_schema = None
                try:
                    if (
                        hasattr(tool, "inputSchema")
                        and tool.inputSchema is not None
                        and isinstance(tool.inputSchema, dict)
                    ):
                        tool_schema = {
                            "type": "function",
                            "function": {
                                "name": tool.name,
                                "description": (
                                    tool.description if hasattr(tool, "description") else None
                                ),
                                "parameters": tool.inputSchema,
                            },
                        }
                except Exception as schema_error:
                    logging.warning(f"Could not extract schema for {tool.name}: {schema_error}")
                    tool_schema = None

                try:
                    tool_obj = Tool(
                        mcp_config=mcp_config,
                        request_options=tool_request_options,
                        tool_schema=tool_schema,
                    )
                    self.register_tool(tool_obj, update=update)
                    registered_tools.append(tool.name)
                except Exception as e:
                    logging.warning(f"Failed to register tool {tool.name}: {e}")

        return registered_tools

    async def load_mcp_config(
        self,
        config_path: str,
        server_names: list[str] | None = None,
        update: bool = False,
        mcp_security: "MCPSecurityConfig | None" = None,
    ) -> dict[str, list[str]]:
        from lionagi.service.connections.mcp_wrapper import (
            MCPConnectionPool,
            MCPSecurityConfig,
        )

        # Explicit config load trusts declared transports by default.
        if mcp_security is None:
            mcp_security = MCPSecurityConfig(allow_commands=True, allow_urls=True)

        MCPConnectionPool.load_config(config_path)

        if server_names is None:
            server_names = list(MCPConnectionPool._configs.keys())
        all_tools = {}
        for server_name in server_names:
            try:
                tools = await self.register_mcp_server(
                    {"server": server_name}, update=update, security=mcp_security
                )
                all_tools[server_name] = tools
                logger.info("Registered %d tools from server '%s'", len(tools), server_name)
            except PermissionError:
                logger.error(
                    "MCP server '%s' was denied by the active MCPSecurityConfig. "
                    "Pass mcp_security=MCPSecurityConfig(allow_commands=True, "
                    "allow_urls=True) (or an allowlist) to load_mcp_config to "
                    "authorize it.",
                    server_name,
                )
                raise
            except Exception as e:
                logger.warning("Failed to register server '%s': %s", server_name, e)
                all_tools[server_name] = []

        return all_tools


async def load_mcp_tools(
    config_path: str | None = None,
    server_names: list[str] | None = None,
    request_options_map: dict[str, dict[str, type]] | None = None,
    update: bool = False,
    mcp_security: "MCPSecurityConfig | None" = None,
) -> list[Tool]:
    from lionagi.service.connections.mcp_wrapper import (
        MCPConnectionPool,
        MCPSecurityConfig,
    )

    manager = ActionManager()

    if mcp_security is None:
        mcp_security = MCPSecurityConfig(allow_commands=True, allow_urls=True)

    if config_path:
        MCPConnectionPool.load_config(config_path)

    if server_names is None and config_path:
        server_names = list(MCPConnectionPool._configs.keys())

    if server_names is None:
        raise ValueError("Either provide server_names or config_path to discover servers")

    for server_name in server_names:
        try:
            request_options = None
            if request_options_map and server_name in request_options_map:
                request_options = request_options_map[server_name]

            tools_registered = await manager.register_mcp_server(
                {"server": server_name},
                request_options=request_options,
                update=update,
                security=mcp_security,
            )
            logger.info("Loaded %d tools from %s", len(tools_registered), server_name)
        except PermissionError:
            logger.error(
                "MCP server '%s' was denied by the active MCPSecurityConfig. "
                "Pass mcp_security=MCPSecurityConfig(allow_commands=True, "
                "allow_urls=True) (or an allowlist) to load_mcp_tools to "
                "authorize it.",
                server_name,
            )
            raise
        except Exception as e:
            logger.warning("Failed to load server '%s': %s", server_name, e)

    return list(manager.registry.values())


__all__ = ["ActionManager", "load_mcp_tools"]

# File: lionagi/protocols/action/manager.py
