from __future__ import annotations

import abc
import asyncio
import warnings
from contextlib import AbstractAsyncContextManager, AsyncExitStack
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal

from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp import ClientSession, StdioServerParameters
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.types import CallToolResult, JSONRPCMessage
from typing_extensions import NotRequired, TypedDict
from .util import MCPUtil

# Allow CAI to talk to newer MCP servers (like SwiftMCP using protocolVersion
# "2025-06-18") without breaking compatibility with older ones. The upstream
# `mcp` client currently only advertises/supports 1 and "2024-11-05", so we
# extend its accepted versions at runtime instead of patching site-packages.
try:  # pragma: no cover - defensive runtime patch
    from mcp.client import session as _mcp_session

    _supported = getattr(_mcp_session, "SUPPORTED_PROTOCOL_VERSIONS", ())
    if "2025-06-18" not in _supported:
        _mcp_session.SUPPORTED_PROTOCOL_VERSIONS = tuple(
            list(_supported) + ["2025-06-18"]
        )
except Exception:
    # If patching fails for any reason, fall back to the library defaults.
    pass

from ..exceptions import UserError
from ..logger import logger

if TYPE_CHECKING:
    # Import Tool for typing only; location is stable across mcp versions.
    from mcp.types import Tool as MCPTool


class MCPServer(abc.ABC):
    """Base class for Model Context Protocol servers."""

    @abc.abstractmethod
    async def connect(self):
        """Connect to the server. For example, this might mean spawning a subprocess or
        opening a network connection. The server is expected to remain connected until
        `cleanup()` is called.
        """
        pass

    @property
    @abc.abstractmethod
    def name(self) -> str:
        """A readable name for the server."""
        pass

    @abc.abstractmethod
    async def cleanup(self):
        """Cleanup the server. For example, this might mean closing a subprocess or
        closing a network connection.
        """
        pass

    @abc.abstractmethod
    async def list_tools(self) -> list[MCPTool]:
        """List the tools available on the server."""
        pass

    @abc.abstractmethod
    async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None) -> CallToolResult:
        """Invoke a tool on the server."""
        pass


class _MCPServerWithClientSession(MCPServer, abc.ABC):
    """Base class for MCP servers that use a `ClientSession` to communicate with the server."""

    def __init__(self, cache_tools_list: bool):
        """
        Args:
            cache_tools_list: Whether to cache the tools list. If `True`, the tools list will be
            cached and only fetched from the server once. If `False`, the tools list will be
            fetched from the server on each call to `list_tools()`. The cache can be invalidated
            by calling `invalidate_tools_cache()`. You should set this to `True` if you know the
            server will not change its tools list, because it can drastically improve latency
            (by avoiding a round-trip to the server every time).
        """
        self.session: ClientSession | None = None
        self.exit_stack: AsyncExitStack = AsyncExitStack()
        self._cleanup_lock: asyncio.Lock = asyncio.Lock()
        self._connect_lock: asyncio.Lock = asyncio.Lock()
        self._call_lock: asyncio.Lock = asyncio.Lock()
        self.cache_tools_list = cache_tools_list

        # The cache is always dirty at startup, so that we fetch tools at least once
        self._cache_dirty = True
        self._tools_list: list[MCPTool] | None = None

    @abc.abstractmethod
    def create_streams(
        self,
    ) -> AbstractAsyncContextManager[
        tuple[
            MemoryObjectReceiveStream[JSONRPCMessage | Exception],
            MemoryObjectSendStream[JSONRPCMessage],
        ]
    ]:
        """Create the streams for the server."""
        pass

    async def __aenter__(self):
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc_value, traceback):
        await self.cleanup()

    def invalidate_tools_cache(self):
        """Invalidate the tools cache."""
        self._cache_dirty = True

    async def connect(self):
        """Connect to the server."""
        if self.session is not None:
            return

        async with self._connect_lock:
            if self.session is not None:
                return

            try:
                transport = await self.exit_stack.enter_async_context(self.create_streams())
                read, write = transport
                session = await self.exit_stack.enter_async_context(ClientSession(read, write))
                await session.initialize()
                self.session = session
            except Exception as e:
                # Only log connection errors at debug level
                error_str = str(e).lower()
                error_type = type(e).__name__
                if (
                    "connection" in error_str
                    or "refused" in error_str
                    or "taskgroup" in error_str
                    or error_type == "ExceptionGroup"
                ):
                    logger.debug(f"Expected connection error during MCP server init: {e}")
                else:
                    logger.error(f"Error initializing MCP server: {e}")
                await self.cleanup()
                raise

    async def list_tools(self) -> list[MCPTool]:
        """List the tools available on the server."""
        if not self.session:
            raise UserError("Server not initialized. Make sure you call `connect()` first.")

        # Return from cache if caching is enabled, we have tools, and the cache is not dirty
        if self.cache_tools_list and not self._cache_dirty and self._tools_list:
            return self._tools_list

        # Reset the cache dirty to False
        self._cache_dirty = False

        # Fetch the tools from the server
        self._tools_list = (await self.session.list_tools()).tools
        return self._tools_list

    async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None) -> CallToolResult:
        """Invoke a tool on the server."""
        if not self.session:
            raise UserError("Server not initialized. Make sure you call `connect()` first.")

        async with self._call_lock:
            if self.session is None:
                await self.connect()

            try:
                return await self.session.call_tool(tool_name, arguments)
            except Exception:
                # Ensure resources are released and mark session stale so callers can reconnect
                await self.cleanup()
                raise

    async def cleanup(self):
        """Cleanup the server."""
        async with self._cleanup_lock:
            try:
                # Suppress async generator warnings during cleanup
                with warnings.catch_warnings():
                    warnings.filterwarnings(
                        "ignore", category=RuntimeWarning, message=".*asynchronous generator.*"
                    )
                    warnings.filterwarnings(
                        "ignore", category=RuntimeWarning, message=".*was never awaited.*"
                    )
                    await self.exit_stack.aclose()
                    self.session = None
            except Exception as e:
                # Only log errors that aren't expected during cleanup
                if "ClosedResourceError" not in str(e) and "async generator" not in str(e).lower():
                    logger.debug(f"Expected cleanup error (can be ignored): {e}")


class MCPServerStdioParams(TypedDict):
    """Mirrors `mcp.client.stdio.StdioServerParameters`, but lets you pass params without another
    import.
    """

    command: str
    """The executable to run to start the server. For example, `python` or `node`."""

    args: NotRequired[list[str]]
    """Command line args to pass to the `command` executable. For example, `['foo.py']` or
    `['server.js', '--port', '8080']`."""

    env: NotRequired[dict[str, str]]
    """The environment variables to set for the server. ."""

    cwd: NotRequired[str | Path]
    """The working directory to use when spawning the process."""

    encoding: NotRequired[str]
    """The text encoding used when sending/receiving messages to the server. Defaults to `utf-8`."""

    encoding_error_handler: NotRequired[Literal["strict", "ignore", "replace"]]
    """The text encoding error handler. Defaults to `strict`.

    See https://docs.python.org/3/library/codecs.html#codec-base-classes for
    explanations of possible values.
    """


class MCPServerStdio(_MCPServerWithClientSession):
    """MCP server implementation that uses the stdio transport. See the [spec]
    (https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#stdio) for
    details.
    """

    def __init__(
        self,
        params: MCPServerStdioParams,
        cache_tools_list: bool = False,
        name: str | None = None,
    ):
        """Create a new MCP server based on the stdio transport.

        Args:
            params: The params that configure the server. This includes the command to run to
                start the server, the args to pass to the command, the environment variables to
                set for the server, the working directory to use when spawning the process, and
                the text encoding used when sending/receiving messages to the server.
            cache_tools_list: Whether to cache the tools list. If `True`, the tools list will be
                cached and only fetched from the server once. If `False`, the tools list will be
                fetched from the server on each call to `list_tools()`. The cache can be
                invalidated by calling `invalidate_tools_cache()`. You should set this to `True`
                if you know the server will not change its tools list, because it can drastically
                improve latency (by avoiding a round-trip to the server every time).
            name: A readable name for the server. If not provided, we'll create one from the
                command.
        """
        super().__init__(cache_tools_list)

        # Build environment more robustly (avoid stdio buffering for Python)
        _env = dict(params.get("env", {}) or {})
        try:
            import os as _os
            cmd_lower = str(params.get("command", "")).lower()
            if "python" in cmd_lower and "PYTHONUNBUFFERED" not in _env:
                _env["PYTHONUNBUFFERED"] = "1"
            # Inherit locale if not provided to avoid encoding issues
            if "LC_ALL" not in _env and "LC_ALL" in _os.environ:
                _env["LC_ALL"] = _os.environ["LC_ALL"]
            # Auto-provide MCP auth token to helper binaries like mcp-proxy
            if "API_ACCESS_TOKEN" not in _env:
                raw_token = MCPUtil._get_default_auth_token_raw()
                if raw_token:
                    _env["API_ACCESS_TOKEN"] = raw_token
        except Exception:
            pass

        self.params = StdioServerParameters(
            command=params["command"],
            args=params.get("args", []),
            env=_env if _env else None,
            cwd=params.get("cwd"),
            encoding=params.get("encoding", "utf-8"),
            # Be tolerant by default to avoid fatal decode errors
            encoding_error_handler=params.get("encoding_error_handler", "replace"),
        )

        self._name = name or f"stdio: {self.params.command}"

    def create_streams(
        self,
    ) -> AbstractAsyncContextManager[
        tuple[
            MemoryObjectReceiveStream[JSONRPCMessage | Exception],
            MemoryObjectSendStream[JSONRPCMessage],
        ]
    ]:
        """Create the streams for the server."""
        return stdio_client(self.params)

    @property
    def name(self) -> str:
        """A readable name for the server."""
        return self._name


class MCPServerSseParams(TypedDict):
    """Mirrors the params in`mcp.client.sse.sse_client`."""

    url: str
    """The URL of the server."""

    headers: NotRequired[dict[str, str]]
    """The headers to send to the server."""

    timeout: NotRequired[float]
    """The timeout for the HTTP request. Defaults to 5 seconds."""

    sse_read_timeout: NotRequired[float]
    """The timeout for the SSE connection, in seconds. Defaults to 5 minutes."""


class MCPServerSse(_MCPServerWithClientSession):
    """MCP server implementation that uses the HTTP with SSE transport. See the [spec]
    (https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#http-with-sse)
    for details.
    """

    def __init__(
        self,
        params: MCPServerSseParams,
        cache_tools_list: bool = False,
        name: str | None = None,
    ):
        """Create a new MCP server based on the HTTP with SSE transport.

        Args:
            params: The params that configure the server. This includes the URL of the server,
                the headers to send to the server, the timeout for the HTTP request, and the
                timeout for the SSE connection.

            cache_tools_list: Whether to cache the tools list. If `True`, the tools list will be
                cached and only fetched from the server once. If `False`, the tools list will be
                fetched from the server on each call to `list_tools()`. The cache can be
                invalidated by calling `invalidate_tools_cache()`. You should set this to `True`
                if you know the server will not change its tools list, because it can drastically
                improve latency (by avoiding a round-trip to the server every time).

            name: A readable name for the server. If not provided, we'll create one from the
                URL.
        """
        super().__init__(cache_tools_list)

        # Ensure sensible defaults and SSE headers
        _headers = MCPUtil.get_default_auth_headers(params.get("headers", {}))
        # Request-side headers (server must still send correct response headers)
        _headers.setdefault("Accept", "text/event-stream")
        _headers.setdefault("Cache-Control", "no-cache")
        _headers.setdefault("Connection", "keep-alive")

        self.params = {
            **params,
            "headers": _headers,
        }
        self._name = name or f"sse: {self.params['url']}"

    def create_streams(
        self,
    ) -> AbstractAsyncContextManager[
        tuple[
            MemoryObjectReceiveStream[JSONRPCMessage | Exception],
            MemoryObjectSendStream[JSONRPCMessage],
        ]
    ]:
        """Create the streams for the server."""
        return sse_client(
            url=self.params["url"],
            headers=self.params.get("headers", None),
            timeout=self.params.get("timeout", 5),
            sse_read_timeout=self.params.get("sse_read_timeout", 60 * 5),
        )

    @property
    def name(self) -> str:
        """A readable name for the server."""
        return self._name

    async def cleanup(self):
        """Cleanup the SSE server with special handling for async generators."""
        import warnings
        import asyncio

        async with self._cleanup_lock:
            try:
                # For SSE servers, we need to handle cleanup more carefully
                with warnings.catch_warnings():
                    warnings.filterwarnings("ignore", category=RuntimeWarning)
                    warnings.filterwarnings("ignore", message=".*asynchronous generator.*")
                    warnings.filterwarnings("ignore", message=".*didn't stop after athrow.*")
                    warnings.filterwarnings("ignore", message=".*cancel scope.*")

                    # Try to close gracefully with a short timeout
                    try:
                        await asyncio.wait_for(self.exit_stack.aclose(), timeout=0.5)
                    except asyncio.TimeoutError:
                        # Expected for SSE connections
                        pass
                    except Exception:
                        # Ignore other cleanup errors for SSE
                        pass

                    self.session = None
            except Exception:
                # Silently ignore all cleanup errors for SSE
                pass

    async def cleanup(self):
        """Cleanup the SSE server with special handling for async generators."""
        async with self._cleanup_lock:
            try:
                # For SSE connections, we need to handle cleanup differently
                # to avoid async generator warnings
                with warnings.catch_warnings():
                    warnings.filterwarnings("ignore", category=RuntimeWarning)
                    warnings.filterwarnings("ignore", message=".*asynchronous generator.*")
                    warnings.filterwarnings("ignore", message=".*was never awaited.*")

                    # Try a quick cleanup with a short timeout
                    try:
                        await asyncio.wait_for(self.exit_stack.aclose(), timeout=0.5)
                    except asyncio.TimeoutError:
                        # Expected for SSE connections
                        pass
                    except Exception:
                        # Ignore any other errors during SSE cleanup
                        pass
                    finally:
                        self.session = None
            except Exception:
                # Silently ignore all errors for SSE cleanup
                pass
