from ..compat import override
from ..entities import ToolCallContext
from . import Tool, ToolSet

from contextlib import AsyncExitStack
from importlib import import_module
from typing import cast


class McpCallTool(Tool):
    """Call an MCP server tool using the MCP client.

    Args:
        uri: Base URI of the MCP server.
        name: Name of the tool to invoke.
        arguments: Arguments to send to the tool.

    Returns:
        Responses returned by the MCP tool invocation.
    """

    _client_params: dict[str, object]
    _call_params: dict[str, object]

    def __init__(
        self,
        *,
        client_params: dict[str, object] | None = None,
        call_params: dict[str, object] | None = None,
    ) -> None:
        super().__init__()
        self.__name__ = "call"
        self._client_params = client_params or {}
        self._call_params = call_params or {}

    async def __call__(
        self,
        uri: str,
        name: str,
        arguments: dict[str, object] | None,
        *,
        context: ToolCallContext,
    ) -> list[object]:
        mcp_module = import_module("mcp")
        client_factory = getattr(mcp_module, "Client")

        assert uri
        assert name

        async with client_factory(uri, **self._client_params) as client:
            result = await client.call_tool(
                name, arguments or {}, **self._call_params
            )
            return cast(list[object], result)


class McpToolSet(ToolSet):
    """Tool set providing MCP client functionality."""

    @override
    def __init__(
        self,
        *,
        exit_stack: AsyncExitStack | None = None,
        namespace: str | None = "mcp",
    ) -> None:
        tools = [McpCallTool()]
        super().__init__(
            exit_stack=exit_stack, namespace=namespace, tools=tools
        )
