"""
Server adapters for different MCP server implementations.

This module contains adapter classes that provide a unified interface
for interacting with different types of MCP servers.
"""

import logging
from abc import ABC, abstractmethod
from typing import Any, Callable, Optional, Protocol, runtime_checkable
from datetime import datetime, timezone

from mcp import ServerResult
from mcp.types import CallToolRequest
from .utils import is_mcp_error_response

logger = logging.getLogger("agnost.analytics.adapters")


@runtime_checkable
class SessionInfo(Protocol):
    """Protocol for session information."""
    session_key: str
    client_name: str


class ServerAdapter(ABC):
    """Abstract base class for MCP server adapters."""

    @abstractmethod
    def get_session_info(self) -> Optional[SessionInfo]:
        """Get session information from the server."""
        pass

    @abstractmethod
    def patch_server(self, analytics_callback: Callable) -> bool:
        """Patch the server to intercept tool calls."""
        pass

    @abstractmethod
    def extract_tools(self) -> list[str]:
        """Extract list of tool names from the server."""
        pass


class FastMCPAdapter(ServerAdapter):
    """Adapter for FastMCP servers."""

    def __init__(self, server: Any):
        self.server = server
        self._mcp_server = server._mcp_server

    def get_session_info(self) -> Optional[SessionInfo]:
        """Get session info from FastMCP server.

        Uses two-path approach matching TypeScript SDK:
        1. Primary: get_client_version() (official MCP SDK method)
        2. Fallback: request_context.session.client_params (legacy path)
        """
        try:
            client_name = 'default'
            session_key = 'fastmcp_default'

            # Try the MCP SDK method first (primary path)
            if hasattr(self._mcp_server, 'get_client_version'):
                try:
                    client_version = self._mcp_server.get_client_version()
                    if client_version and hasattr(client_version, 'name') and client_version.name:
                        client_name = client_version.name
                        logger.debug(f"Got client name from get_client_version: {client_name}")
                except Exception as e:
                    logger.debug(f"get_client_version failed: {e}")

            # Fallback to request_context path for compatibility
            if client_name == 'default':
                if hasattr(self._mcp_server, 'request_context') and self._mcp_server.request_context:
                    if hasattr(self._mcp_server.request_context, 'session') and self._mcp_server.request_context.session:
                        session = self._mcp_server.request_context.session
                        session_key = hex(id(session))
                        if hasattr(session, 'client_params') and session.client_params:
                            if hasattr(session.client_params, 'clientInfo') and session.client_params.clientInfo:
                                if hasattr(session.client_params.clientInfo, 'name'):
                                    client_name = session.client_params.clientInfo.name
                                    logger.debug(f"Got client name from request_context: {client_name}")

            return type('SessionInfo', (), {
                'session_key': session_key,
                'client_name': client_name
            })()
        except Exception as e:
            logger.debug(f"FastMCP session info error: {e}")
            return None

    def patch_server(self, analytics_callback: Callable) -> bool:
        """Patch FastMCP server by replacing the CallToolRequest handler."""
        try:
            # Patch the CallToolRequest handler to intercept tool calls
            # We can't patch FastMCP.call_tool directly because the handler
            # was already registered with a closure that captured the original method
            mcp_server = self._mcp_server
            original_handler = mcp_server.request_handlers.get(CallToolRequest)

            if not original_handler:
                logger.error("No CallToolRequest handler found")
                return False

            async def wrapped_handler(request: CallToolRequest):
                tool_name = request.params.name
                arguments = request.params.arguments or {}

                start_time = datetime.now(timezone.utc)
                success = True
                result = None

                try:
                    exec_start = datetime.now(timezone.utc)
                    result = await original_handler(request)
                    exec_end = datetime.now(timezone.utc)
                    exec_time = int((exec_end - exec_start).total_seconds() * 1000)

                    is_error, error_message = is_mcp_error_response(result)
                    if is_error:
                        success = False
                        logger.error(f"Tool {tool_name} returned error: {error_message}")

                except Exception as e:
                    success = False
                    exec_time = int((datetime.now(timezone.utc) - start_time).total_seconds() * 1000)
                    logger.error(f"Error calling tool {tool_name}: {e}")
                    raise
                finally:
                    analytics_callback(tool_name, arguments, exec_time, success, result, start_time)

                return result

            # Replace the handler
            mcp_server.request_handlers[CallToolRequest] = wrapped_handler

            logger.debug("FastMCP server patched successfully")
            return True
        except Exception as e:
            logger.error(f"Failed to patch FastMCP server: {e}")
            return False

    def extract_tools(self) -> list[str]:
        """Extract list of tool names from FastMCP server.

        For FastMCP, tools are available immediately after registration via @mcp.tool() decorators.
        This method will be called during session creation to capture the initial tool list.
        """
        try:
            # Access the tool manager's tool registry
            if hasattr(self.server, '_tool_manager'):
                tool_manager = self.server._tool_manager

                # FastMCP stores tools in _tools dict after registration
                if hasattr(tool_manager, '_tools') and tool_manager._tools:
                    tools = list(tool_manager._tools.keys())
                    logger.info(f"Extracted {len(tools)} tools from FastMCP: {', '.join(tools)}")
                    return tools
                else:
                    logger.warning("FastMCP tool manager _tools is empty - tools may not be registered yet")
                    return []

            logger.warning("FastMCP server missing _tool_manager attribute")
            return []
        except Exception as e:
            logger.error(f"Failed to extract tools from FastMCP server: {e}")
            return []


class LowLevelMCPAdapter(ServerAdapter):
    """Adapter for low-level MCP servers."""

    def __init__(self, server: Any):
        self.server = server._mcp_server if hasattr(server, '_mcp_server') else server

    def get_session_info(self) -> Optional[SessionInfo]:
        """Get session info from low-level MCP server.

        Uses two-path approach matching TypeScript SDK:
        1. Primary: get_client_version() (official MCP SDK method)
        2. Fallback: request_context.session.client_params (legacy path)
        """
        try:
            client_name = 'default'
            session_key = 'lowlevel_default'

            # Try the MCP SDK method first (primary path)
            if hasattr(self.server, 'get_client_version'):
                try:
                    client_version = self.server.get_client_version()
                    if client_version and hasattr(client_version, 'name') and client_version.name:
                        client_name = client_version.name
                        logger.debug(f"Got client name from get_client_version: {client_name}")
                except Exception as e:
                    logger.debug(f"get_client_version failed: {e}")

            # Fallback to request_context path for compatibility
            if client_name == 'default':
                if hasattr(self.server, 'request_context') and self.server.request_context:
                    if hasattr(self.server.request_context, 'session') and self.server.request_context.session:
                        session = self.server.request_context.session
                        session_key = hex(id(session))
                        if hasattr(session, 'client_params') and session.client_params:
                            if hasattr(session.client_params, 'clientInfo') and session.client_params.clientInfo:
                                if hasattr(session.client_params.clientInfo, 'name'):
                                    client_name = session.client_params.clientInfo.name
                                    logger.debug(f"Got client name from request_context: {client_name}")

            return type('SessionInfo', (), {
                'session_key': session_key,
                'client_name': client_name
            })()
        except Exception as e:
            logger.debug(f"LowLevel MCP session info error: {e}")
            return None

    def patch_server(self, analytics_callback: Callable) -> bool:
        """Patch low-level MCP server request handlers."""
        try:
            original_handler = self.server.request_handlers.get(CallToolRequest)
            if not original_handler:
                logger.error("No CallToolRequest handler found")
                return False

            async def wrapped_handler(request: CallToolRequest) -> ServerResult:
                tool_name = request.params.name
                arguments = request.params.arguments or {}

                start_time = datetime.now(timezone.utc)
                success = True
                result = None

                try:
                    exec_start = datetime.now(timezone.utc)
                    result = await original_handler(request)
                    exec_end = datetime.now(timezone.utc)
                    exec_time = int((exec_end - exec_start).total_seconds() * 1000)

                    is_error, error_message = is_mcp_error_response(result)
                    if is_error:
                        success = False
                        logger.error(f"Tool {tool_name} returned error: {error_message}")

                except Exception as e:
                    success = False
                    exec_time = int((datetime.now(timezone.utc) - start_time).total_seconds() * 1000)
                    logger.error(f"Error calling tool {tool_name}: {e}")
                    raise
                finally:
                    analytics_callback(tool_name, arguments, exec_time, success, result, start_time)

                return result

            self.server.request_handlers[CallToolRequest] = wrapped_handler
            logger.debug("Low-level MCP server patched successfully")
            return True
        except Exception as e:
            logger.error(f"Failed to patch low-level MCP server: {e}")
            return False

    def extract_tools(self) -> list[str]:
        """Extract list of tool names from low-level MCP server.

        Note: For low-level MCP servers, tool extraction works best when tools are
        properly registered and the server is fully initialized. In test/mock contexts,
        this may return an empty list.

        In production, tools will be extracted when the server processes the first
        tools/list request from a client.
        """
        try:
            # Access the _tool_cache which stores tool definitions after they're cached
            if hasattr(self.server, '_tool_cache') and self.server._tool_cache:
                tools = list(self.server._tool_cache.keys())
                logger.info(f"Extracted {len(tools)} tools from low-level MCP: {', '.join(tools)}")
                return tools

            # For low-level MCP servers in test/early initialization,
            # the cache may be empty. In production use, tools are populated
            # when the server handles a real tools/list request from a client.
            logger.debug("No tools found in _tool_cache (server may not be fully initialized)")
            return []

        except Exception as e:
            logger.warning(f"Failed to extract tools from low-level MCP server: {e}")
            return []