"""
A derived client session for the MCP Agent framework.
It adds logging and supports sampling requests.
"""

from datetime import timedelta
from typing import TYPE_CHECKING, Any

from mcp import ClientSession, ServerNotification
from mcp.shared.context import RequestContext
from mcp.shared.message import MessageMetadata
from mcp.shared.session import (
    ProgressFnT,
    ReceiveResultT,
)
from mcp.types import (
    CallToolRequest,
    CallToolRequestParams,
    CallToolResult,
    ClientRequest,
    EmptyResult,
    GetPromptRequest,
    GetPromptRequestParams,
    GetPromptResult,
    Implementation,
    ListRootsResult,
    PingRequest,
    ReadResourceRequest,
    ReadResourceRequestParams,
    ReadResourceResult,
    Root,
    SamplingCapability,
    SamplingToolsCapability,
    ToolListChangedNotification,
)
from pydantic import AnyUrl, FileUrl

from fast_agent.context_dependent import ContextDependent
from fast_agent.core.logging.logger import get_logger
from fast_agent.mcp.helpers.server_config_helpers import get_server_config
from fast_agent.mcp.sampling import sample

if TYPE_CHECKING:
    from fast_agent.config import MCPServerSettings
    from fast_agent.mcp.transport_tracking import TransportChannelMetrics

logger = get_logger(__name__)


async def list_roots(context: RequestContext[ClientSession, None]) -> ListRootsResult:
    """List roots callback that will be called by the MCP library."""

    if server_config := get_server_config(context.session):
        if server_config.roots:
            roots = [
                Root(
                    uri=FileUrl(
                        root.server_uri_alias or root.uri,
                    ),
                    name=root.name,
                )
                for root in server_config.roots
            ]
            return ListRootsResult(roots=roots)

    return ListRootsResult(roots=[])


class MCPAgentClientSession(ClientSession, ContextDependent):
    """
    MCP Agent framework acts as a client to the servers providing tools/resources/prompts for the agent workloads.
    This is a simple client session for those server connections, and supports
        - handling sampling requests
        - notifications
        - MCP root configuration

    Developers can extend this class to add more custom functionality as needed
    """

    def __init__(self, read_stream, write_stream, read_timeout=None, **kwargs) -> None:
        # Extract server_name if provided in kwargs
        from importlib.metadata import version

        self.session_server_name = kwargs.pop("server_name", None)
        # Extract the notification callbacks if provided
        self._tool_list_changed_callback = kwargs.pop("tool_list_changed_callback", None)
        # Extract server_config if provided
        self.server_config: MCPServerSettings | None = kwargs.pop("server_config", None)
        # Extract agent_model if provided (for auto_sampling fallback)
        self.agent_model: str | None = kwargs.pop("agent_model", None)
        # Extract agent_name if provided
        self.agent_name: str | None = kwargs.pop("agent_name", None)
        # Extract api_key if provided
        self.api_key: str | None = kwargs.pop("api_key", None)
        # Extract custom elicitation handler if provided
        custom_elicitation_handler = kwargs.pop("elicitation_handler", None)
        # Extract optional context for ContextDependent mixin without passing it to ClientSession
        self._context = kwargs.pop("context", None)
        # Extract transport metrics tracker if provided
        self._transport_metrics: TransportChannelMetrics | None = kwargs.pop(
            "transport_metrics", None
        )

        # Track the effective elicitation mode for diagnostics
        self.effective_elicitation_mode: str | None = "none"

        fast_agent_version = version("fast-agent-mcp") or "dev"
        fast_agent: Implementation = Implementation(name="fast-agent-mcp", version=fast_agent_version)
        if self.server_config and self.server_config.implementation:
            fast_agent = self.server_config.implementation

        # Only register callbacks if the server_config has the relevant settings
        list_roots_cb = list_roots if (self.server_config and self.server_config.roots) else None

        # Register sampling callback if either:
        # 1. Sampling is explicitly configured, OR
        # 2. Application-level auto_sampling is enabled
        sampling_cb = None
        if self.server_config and self.server_config.sampling:
            # Explicit sampling configuration
            sampling_cb = sample
        elif self._should_enable_auto_sampling():
            # Auto-sampling enabled at application level
            sampling_cb = sample

        # Use custom elicitation handler if provided, otherwise resolve using factory
        if custom_elicitation_handler is not None:
            elicitation_handler = custom_elicitation_handler
        else:
            # Try to resolve using factory
            elicitation_handler = None
            try:
                from fast_agent.agents.agent_types import AgentConfig
                from fast_agent.context import get_current_context
                from fast_agent.mcp.elicitation_factory import resolve_elicitation_handler

                context = get_current_context()
                if context and context.config:
                    # Create a minimal agent config for the factory
                    agent_config = AgentConfig(
                        name=self.agent_name or "unknown",
                        model=self.agent_model or "unknown",
                        elicitation_handler=None,
                    )
                    elicitation_handler = resolve_elicitation_handler(
                        agent_config, context.config, self.server_config
                    )
            except Exception:
                # If factory resolution fails, we'll use default fallback
                pass

            # Fallback to forms handler only if factory resolution wasn't attempted
            if elicitation_handler is None and not self.server_config:
                from fast_agent.mcp.elicitation_handlers import forms_elicitation_handler

                elicitation_handler = forms_elicitation_handler

        # Determine effective elicitation mode for diagnostics
        if self.server_config and getattr(self.server_config, "elicitation", None):
            self.effective_elicitation_mode = self.server_config.elicitation.mode or "forms"
        elif elicitation_handler is not None:
            # Use global config if available to distinguish auto-cancel
            try:
                from fast_agent.context import get_current_context

                context = get_current_context()
                mode = None
                if context and getattr(context, "config", None):
                    elicitation_cfg = getattr(context.config, "elicitation", None)
                    if isinstance(elicitation_cfg, dict):
                        mode = elicitation_cfg.get("mode")
                    else:
                        mode = getattr(elicitation_cfg, "mode", None)
                self.effective_elicitation_mode = (mode or "forms").lower()
            except Exception:
                self.effective_elicitation_mode = "forms"
        else:
            self.effective_elicitation_mode = "none"

        # Pop parameters we're explicitly setting to avoid duplicates
        kwargs.pop("list_roots_callback", None)
        kwargs.pop("sampling_callback", None)
        kwargs.pop("sampling_capabilities", None)
        kwargs.pop("client_info", None)
        kwargs.pop("elicitation_callback", None)

        # Create sampling capabilities with tools support when sampling is enabled
        sampling_caps = None
        if sampling_cb is not None:
            # Advertise full sampling capability including tools support
            sampling_caps = SamplingCapability(
                tools=SamplingToolsCapability()
            )

        super().__init__(
            read_stream,
            write_stream,
            read_timeout,
            **kwargs,
            list_roots_callback=list_roots_cb,
            sampling_callback=sampling_cb,
            sampling_capabilities=sampling_caps,
            client_info=fast_agent,
            elicitation_callback=elicitation_handler,
        )

    def _should_enable_auto_sampling(self) -> bool:
        """Check if auto_sampling is enabled at the application level."""
        try:
            from fast_agent.context import get_current_context

            context = get_current_context()
            if context and context.config:
                return getattr(context.config, "auto_sampling", True)
        except Exception:
            pass
        return True  # Default to True if can't access config

    async def send_request(
        self,
        request: ClientRequest,
        result_type: type[ReceiveResultT],
        request_read_timeout_seconds: timedelta | None = None,
        metadata: MessageMetadata | None = None,
        progress_callback: ProgressFnT | None = None,
    ) -> ReceiveResultT:
        logger.debug("send_request: request=", data=request.model_dump())
        request_id = getattr(self, "_request_id", None)
        is_ping_request = self._is_ping_request(request)
        if (
            is_ping_request
            and request_id is not None
            and self._transport_metrics is not None
        ):
            self._transport_metrics.register_ping_request(request_id)
        try:
            result = await super().send_request(
                request=request,
                result_type=result_type,
                request_read_timeout_seconds=request_read_timeout_seconds,
                metadata=metadata,
                progress_callback=progress_callback,
            )
            logger.debug(
                "send_request: response=",
                data=result.model_dump() if result is not None else "no response returned",
            )
            self._attach_transport_channel(request_id, result)
            if (
                is_ping_request
                and request_id is not None
                and self._transport_metrics is not None
            ):
                self._transport_metrics.discard_ping_request(request_id)
            return result
        except Exception as e:
            if (
                is_ping_request
                and request_id is not None
                and self._transport_metrics is not None
            ):
                self._transport_metrics.discard_ping_request(request_id)
            from anyio import ClosedResourceError

            from fast_agent.core.exceptions import ServerSessionTerminatedError

            # Check for session terminated error (404 from server)
            if self._is_session_terminated_error(e):
                raise ServerSessionTerminatedError(
                    server_name=self.session_server_name or "unknown",
                    details="Server returned 404 - session may have expired due to server restart",
                ) from e

            # Handle connection closure errors (transport closed)
            if isinstance(e, ClosedResourceError):
                from fast_agent.ui import console

                console.console.print(
                    f"[dim red]MCP server {self.session_server_name} offline[/dim red]"
                )
                raise ConnectionError(f"MCP server {self.session_server_name} offline") from e

            logger.error(f"send_request failed: {str(e)}")
            raise

    @staticmethod
    def _is_ping_request(request: ClientRequest) -> bool:
        root = getattr(request, "root", None)
        method = getattr(root, "method", None)
        if not isinstance(method, str):
            return False
        method_lower = method.lower()
        return method_lower == "ping" or method_lower.endswith("/ping") or method_lower.endswith(".ping")

    def _is_session_terminated_error(self, exc: Exception) -> bool:
        """Check if exception is a session terminated error (code 32600 from 404)."""
        from mcp.shared.exceptions import McpError

        from fast_agent.core.exceptions import ServerSessionTerminatedError

        if isinstance(exc, McpError):
            error_data = getattr(exc, "error", None)
            if error_data:
                code = getattr(error_data, "code", None)
                if code == ServerSessionTerminatedError.SESSION_TERMINATED_CODE:
                    return True
        return False

    def _attach_transport_channel(self, request_id, result) -> None:
        if self._transport_metrics is None or request_id is None or result is None:
            return
        channel = self._transport_metrics.consume_response_channel(request_id)
        if not channel:
            return
        try:
            setattr(result, "transport_channel", channel)
        except Exception:
            # If result cannot be mutated, ignore silently
            pass

    async def _received_notification(self, notification: ServerNotification) -> None:
        """
        Can be overridden by subclasses to handle a notification without needing
        to listen on the message stream.
        """
        logger.debug(
            "_received_notification: notification=",
            data=notification.model_dump(),
        )

        # Call parent notification handler first
        await super()._received_notification(notification)

        # Then process our specific notification types
        match notification.root:
            case ToolListChangedNotification():
                # Simple notification handling - just call the callback if it exists
                if self._tool_list_changed_callback and self.session_server_name:
                    logger.info(
                        f"Tool list changed for server '{self.session_server_name}', triggering callback"
                    )
                    # Use asyncio.create_task to prevent blocking the notification handler
                    import asyncio

                    asyncio.create_task(
                        self._handle_tool_list_change_callback(self.session_server_name)
                    )
                else:
                    logger.debug(
                        f"Tool list changed for server '{self.session_server_name}' but no callback registered"
                    )

        return None

    async def _handle_tool_list_change_callback(self, server_name: str) -> None:
        """
        Helper method to handle tool list change callback in a separate task
        to prevent blocking the notification handler
        """
        try:
            await self._tool_list_changed_callback(server_name)
        except Exception as e:
            logger.error(f"Error in tool list changed callback: {e}")

    async def call_tool(
        self,
        name: str,
        arguments: dict[str, Any] | None = None,
        read_timeout_seconds: timedelta | None = None,
        progress_callback: ProgressFnT | None = None,
        *,
        meta: dict[str, Any] | None = None,
    ) -> CallToolResult:
        """Call a tool with optional metadata and progress callback support.

        Always uses our overridden send_request to ensure session terminated errors
        are properly detected and converted to ServerSessionTerminatedError.
        """
        from mcp.types import RequestParams

        # Always create request ourselves to ensure we go through our send_request override
        # This is critical for session terminated detection to work
        _meta: RequestParams.Meta | None = None
        if meta is not None:
            _meta = RequestParams.Meta(**meta)

        # ty doesn't recognize _meta from pydantic alias - this matches SDK pattern
        params = CallToolRequestParams(name=name, arguments=arguments, _meta=_meta)  # ty: ignore[unknown-argument]

        request = CallToolRequest(method="tools/call", params=params)
        return await self.send_request(
            ClientRequest(request),
            CallToolResult,
            request_read_timeout_seconds=read_timeout_seconds,
            progress_callback=progress_callback,
        )

    async def ping(self, read_timeout_seconds: timedelta | None = None) -> EmptyResult:
        """Send a ping request to check server liveness."""
        request = PingRequest(method="ping")
        return await self.send_request(
            ClientRequest(request),
            EmptyResult,
            request_read_timeout_seconds=read_timeout_seconds,
        )

    async def read_resource(
        self, uri: AnyUrl | str, _meta: dict | None = None, **kwargs
    ) -> ReadResourceResult:
        """Read a resource with optional metadata support.

        Always uses our overridden send_request to ensure session terminated errors
        are properly detected and converted to ServerSessionTerminatedError.
        """
        from mcp.types import RequestParams

        # Convert str to AnyUrl if needed
        uri_obj: AnyUrl = uri if isinstance(uri, AnyUrl) else AnyUrl(uri)

        # Always create request ourselves to ensure we go through our send_request override
        params = ReadResourceRequestParams(uri=uri_obj)

        if _meta:
            # Safe merge - preserve existing meta fields like progressToken
            existing_meta = kwargs.get("meta")
            if existing_meta:
                meta_dict = (
                    existing_meta.model_dump() if hasattr(existing_meta, "model_dump") else {}
                )
                meta_dict.update(_meta)
                meta_obj = RequestParams.Meta(**meta_dict)
            else:
                meta_obj = RequestParams.Meta(**_meta)
            params = ReadResourceRequestParams(uri=uri_obj, meta=meta_obj)

        request = ReadResourceRequest(method="resources/read", params=params)
        return await self.send_request(ClientRequest(request), ReadResourceResult)

    async def get_prompt(
        self, name: str, arguments: dict | None = None, _meta: dict | None = None, **kwargs
    ) -> GetPromptResult:
        """Get a prompt with optional metadata support.

        Always uses our overridden send_request to ensure session terminated errors
        are properly detected and converted to ServerSessionTerminatedError.
        """
        from mcp.types import RequestParams

        # Always create request ourselves to ensure we go through our send_request override
        params = GetPromptRequestParams(name=name, arguments=arguments)

        if _meta:
            # Safe merge - preserve existing meta fields like progressToken
            existing_meta = kwargs.get("meta")
            if existing_meta:
                meta_dict = (
                    existing_meta.model_dump() if hasattr(existing_meta, "model_dump") else {}
                )
                meta_dict.update(_meta)
                meta_obj = RequestParams.Meta(**meta_dict)
            else:
                meta_obj = RequestParams.Meta(**_meta)
            params = GetPromptRequestParams(name=name, arguments=arguments, meta=meta_obj)

        request = GetPromptRequest(method="prompts/get", params=params)
        return await self.send_request(ClientRequest(request), GetPromptResult)
