"""Core module for the MCPAdapt library.

This module contains the core functionality for the MCPAdapt library. It provides the
basic interfaces and classes for adapting tools from MCP to the desired Agent framework.
"""

import asyncio
import threading
from abc import ABC, abstractmethod
from contextlib import asynccontextmanager
from functools import partial
from typing import Any, Callable, Coroutine

import mcp
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client


class ToolAdapter(ABC):
    """A basic interface for adapting tools from MCP to the desired Agent framework."""

    @abstractmethod
    def adapt(
        self,
        func: Callable[[dict | None], mcp.types.CallToolResult],
        mcp_tool: mcp.types.Tool,
    ):
        """Adapt a single tool from MCP to the desired Agent framework.

        The MCP protocol will provide a name, description and inputSchema in JSON Schema
        format. This needs to be adapted to the desired Agent framework.

        Note that the function is synchronous (not a coroutine) you can use
        :meth:`ToolAdapter.async_adapt` if you need to use the tool asynchronously.

        Args:
            func: The function to be called.
            mcp_tool: The tool to adapt.

        Returns:
            The adapted tool.
        """
        pass

    def async_adapt(
        self,
        afunc: Callable[[dict | None], Coroutine[Any, Any, mcp.types.CallToolResult]],
        mcp_tool: mcp.types.Tool,
    ):
        """Adapt a single tool from MCP to the desired Agent framework.

        The MCP protocol will provide a name, description and inputSchema in JSON Schema
        format. This needs to be adapted to the desired Agent framework.

        Note that the function is asynchronous (a coroutine) you can use
        :meth:`ToolAdapter.adapt` if you need to use the tool synchronously.

        Args:
            afunc: The coroutine to be called.
            mcp_tool: The tool to adapt.

        Returns:
            The adapted tool.
        """
        raise NotImplementedError(
            "Async adaptation is not supported for this Agent framework."
        )


@asynccontextmanager
async def mcptools(
    serverparams: StdioServerParameters,
) -> tuple[ClientSession, list[mcp.types.Tool]]:
    """Async context manager that yields tools from an MCP server.

    Note: the session can be then used to call tools on the MCP server but it's async.
    Use MCPAdapt instead if you need to use the tools synchronously.

    Args:
        serverparams: Parameters to run the MCP server.

    Yields:
        A list of tools available on the MCP server.

    Usage:
    >>> async with mcptools(StdioServerParameters(command="uv", args=["run", "src/echo.py"])) as (session, tools):
    >>>     print(tools)
    """
    async with stdio_client(serverparams) as (read, write):
        async with ClientSession(read, write) as session:
            # Initialize the connection and get the tools from the mcp server
            await session.initialize()
            tools = await session.list_tools()
            yield session, tools.tools


class MCPAdapt:
    """The main class for adapting MCP tools to the desired Agent framework.

    This class can be used either as a sync or async context manager.

    If running synchronously, it will run the MCP server in a separate thread and take
    care of making the tools synchronous without blocking the server.

    If running asynchronously, it will use the async context manager and return async
    tools.

    Dependening on what your Agent framework supports choose the approriate method. If
    async is supported it is recommended.

    Important Note: adapters need to implement the async_adapt method to support async
    tools.

    Usage:
    >>> # sync usage
    >>> with MCPAdapt(StdioServerParameters(command="uv", args=["run", "src/echo.py"]), SmolAgentAdapter()) as tools:
    >>>     print(tools)

    >>> # async usage
    >>> async with MCPAdapt(StdioServerParameters(command="uv", args=["run", "src/echo.py"]), SmolAgentAdapter()) as tools:
    >>>     print(tools)
    """

    def __init__(self, serverparams: StdioServerParameters, adapter: ToolAdapter):
        # attributes we receive from the user.
        self.serverparams = serverparams
        self.adapter = adapter

        # session and tools get set by the async loop during initialization.
        self.session: ClientSession = None
        self.mcp_tools: list[mcp.types.Tool] = None

        # all attributes used to manage the async loop and separate thread.
        self.loop = asyncio.new_event_loop()
        self.task = None
        self.ready = threading.Event()
        self.thread = threading.Thread(target=self._run_loop, daemon=True)

        # start the loop in a separate thread and wait till ready synchronously.
        self.thread.start()
        self.ready.wait()

    def _run_loop(self):
        """Runs the event loop in a separate thread (for synchronous usage)."""
        asyncio.set_event_loop(self.loop)

        async def setup():
            async with mcptools(self.serverparams) as (session, tools):
                self.session, self.mcp_tools = session, tools
                self.ready.set()  # Signal initialization is complete
                await asyncio.Event().wait()  # Keep session alive until stopped

        self.task = self.loop.create_task(setup())
        try:
            self.loop.run_until_complete(self.task)
        except asyncio.CancelledError:
            pass

    def tools(self) -> list[Any]:
        """Returns the tools from the MCP server adapted to the desired Agent framework.

        This is what is yielded if used as a context manager otherwise you can access it
        directly via this method.

        An equivalent async method is available if your Agent framework supports it:
        see :meth:`atools`.

        """
        if not self.session:
            raise RuntimeError("Session not initialized")

        def _sync_call_tool(
            name: str, arguments: dict | None = None
        ) -> mcp.types.CallToolResult:
            return asyncio.run_coroutine_threadsafe(
                self.session.call_tool(name, arguments), self.loop
            ).result()

        return [
            self.adapter.adapt(partial(_sync_call_tool, tool.name), tool)
            for tool in self.mcp_tools
        ]

    def close(self):
        """Clean up resources and stop the client."""
        if self.task and not self.task.done():
            self.loop.call_soon_threadsafe(self.task.cancel)
        self.thread.join()  # will wait until the task is cancelled to join thread (as it's blocked Event().wait())
        self.loop.close()  # we won't be using the loop anymore we can safely close it

    def __enter__(self):
        return self.tools()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    # -- add support for async context manager as well if the agent framework supports it.
    def atools(self) -> list[Any]:
        """Returns the tools from the MCP server adapted to the desired Agent framework.

        This is what is yielded if used as an async context manager otherwise you can
        access it directly via this method.

        An equivalent async method is available if your Agent framework supports it:
        see :meth:`atools`.
        """
        return [
            self.adapter.async_adapt(partial(self.session.call_tool, tool.name), tool)
            for tool in self.mcp_tools
        ]

    async def __aenter__(self) -> list[Any]:
        self._ctxmanager = mcptools(self.serverparams)
        self.session, self.tools = await self._ctxmanager.__aenter__()
        return self.atools()

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self._ctxmanager.__aexit__(exc_type, exc_val, exc_tb)


if __name__ == "__main__":

    class DummyAdapter(ToolAdapter):
        def adapt(
            self,
            func: Callable[[dict | None], mcp.types.CallToolResult],
            mcp_tool: mcp.types.Tool,
        ):
            return func

        def async_adapt(
            self,
            afunc: Callable[
                [dict | None], Coroutine[Any, Any, mcp.types.CallToolResult]
            ],
            mcp_tool: mcp.types.Tool,
        ):
            return afunc

    with MCPAdapt(
        StdioServerParameters(command="uv", args=["run", "src/echo.py"]),
        DummyAdapter(),
    ) as smolagents_tools:
        print(smolagents_tools)
        print(smolagents_tools[0].forward({"text": "hello"}))

    async def main():
        async with MCPAdapt(
            StdioServerParameters(command="uv", args=["run", "src/echo.py"]),
            DummyAdapter(),
        ) as smolagents_tools:
            print(smolagents_tools)
            print(smolagents_tools[0].forward({"text": "hello"}))

    asyncio.run(main())
