# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import asyncio
from collections import deque
from contextlib import AsyncExitStack
from datetime import timedelta
import functools
import hashlib
import json
import logging
import sys
import threading
from typing import Any
from typing import Dict
from typing import Optional
from typing import Protocol
from typing import runtime_checkable
from typing import TextIO
from typing import Union

from mcp import ClientSession
from mcp import StdioServerParameters
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.client.streamable_http import create_mcp_http_client
from mcp.client.streamable_http import McpHttpClientFactory
from mcp.client.streamable_http import streamablehttp_client
from pydantic import BaseModel
from pydantic import ConfigDict

from .session_context import SessionContext

logger = logging.getLogger('google_adk.' + __name__)


def _has_cancelled_error_context(exc: BaseException) -> bool:
  """Returns True if `exc` is/was caused by `asyncio.CancelledError`.

  Cancellation can be translated into other exceptions during teardown (e.g.
  connection errors) while still retaining the original cancellation in an
  exception's context chain.
  """

  seen: set[int] = set()
  queue = deque([exc])
  while queue:
    current = queue.popleft()
    if id(current) in seen:
      continue
    seen.add(id(current))
    if isinstance(current, asyncio.CancelledError):
      return True
    if current.__cause__ is not None:
      queue.append(current.__cause__)
    if current.__context__ is not None:
      queue.append(current.__context__)
  return False


class StdioConnectionParams(BaseModel):
  """Parameters for the MCP Stdio connection.

  Attributes:
      server_params: Parameters for the MCP Stdio server.
      timeout: Timeout in seconds for establishing the connection to the MCP
        stdio server.
  """

  server_params: StdioServerParameters
  timeout: float = 5.0


class SseConnectionParams(BaseModel):
  """Parameters for the MCP SSE connection.

  See MCP SSE Client documentation for more details.
  https://github.com/modelcontextprotocol/python-sdk/blob/main/src/mcp/client/sse.py

  Attributes:
      url: URL for the MCP SSE server.
      headers: Headers for the MCP SSE connection.
      timeout: Timeout in seconds for establishing the connection to the MCP SSE
        server.
      sse_read_timeout: Timeout in seconds for reading data from the MCP SSE
        server.
  """

  url: str
  headers: dict[str, Any] | None = None
  timeout: float = 5.0
  sse_read_timeout: float = 60 * 5.0


@runtime_checkable
class CheckableMcpHttpClientFactory(McpHttpClientFactory, Protocol):
  pass


class StreamableHTTPConnectionParams(BaseModel):
  """Parameters for the MCP Streamable HTTP connection.

  See MCP Streamable HTTP Client documentation for more details.
  https://github.com/modelcontextprotocol/python-sdk/blob/main/src/mcp/client/streamable_http.py

  Attributes:
      url: URL for the MCP Streamable HTTP server.
      headers: Headers for the MCP Streamable HTTP connection.
      timeout: Timeout in seconds for establishing the connection to the MCP
        Streamable HTTP server.
      sse_read_timeout: Timeout in seconds for reading data from the MCP
        Streamable HTTP server.
      terminate_on_close: Whether to terminate the MCP Streamable HTTP server
        when the connection is closed.
      httpx_client_factory: Factory function to create a custom HTTPX client. If
        not provided, a default factory will be used.
  """

  model_config = ConfigDict(arbitrary_types_allowed=True)

  url: str
  headers: dict[str, Any] | None = None
  timeout: float = 5.0
  sse_read_timeout: float = 60 * 5.0
  terminate_on_close: bool = True
  httpx_client_factory: CheckableMcpHttpClientFactory = create_mcp_http_client


def retry_on_errors(func):
  """Decorator to automatically retry action when MCP session errors occur.

  When MCP session errors occur, the decorator will automatically retry the
  action once. The create_session method will handle creating a new session
  if the old one was disconnected.

  Cancellation is not retried and must be allowed to propagate. In async
  runtimes, cancellation may surface as `asyncio.CancelledError` or as another
  exception while the task is cancelling.

  Args:
      func: The function to decorate.

  Returns:
      The decorated function.
  """

  @functools.wraps(func)  # Preserves original function metadata
  async def wrapper(self, *args, **kwargs):
    try:
      return await func(self, *args, **kwargs)
    except Exception as e:
      task = asyncio.current_task()
      if task is not None:
        cancelling = getattr(task, 'cancelling', None)
        if cancelling is not None and cancelling() > 0:
          raise
      if _has_cancelled_error_context(e):
        raise
      # If an error is thrown, we will retry the function to reconnect to the
      # server. create_session will handle detecting and replacing disconnected
      # sessions.
      logger.info('Retrying %s due to error: %s', func.__name__, e)
      return await func(self, *args, **kwargs)

  return wrapper


class MCPSessionManager:
  """Manages MCP client sessions.

  This class provides methods for creating and initializing MCP client sessions,
  handling different connection parameters (Stdio and SSE) and supporting
  session pooling based on authentication headers.
  """

  def __init__(
      self,
      connection_params: Union[
          StdioServerParameters,
          StdioConnectionParams,
          SseConnectionParams,
          StreamableHTTPConnectionParams,
      ],
      errlog: TextIO = sys.stderr,
  ):
    """Initializes the MCP session manager.

    Args:
        connection_params: Parameters for the MCP connection (Stdio, SSE or
          Streamable HTTP). Stdio by default also has a 5s read timeout as other
          parameters but it's not configurable for now.
        errlog: (Optional) TextIO stream for error logging. Use only for
          initializing a local stdio MCP session.
    """
    if isinstance(connection_params, StdioServerParameters):
      # So far timeout is not configurable. Given MCP is still evolving, we
      # would expect stdio_client to evolve to accept timeout parameter like
      # other client.
      logger.warning(
          'StdioServerParameters is not recommended. Please use'
          ' StdioConnectionParams.'
      )
      self._connection_params = StdioConnectionParams(
          server_params=connection_params,
          timeout=5,
      )
    else:
      self._connection_params = connection_params
    self._errlog = errlog

    # Session pool: maps session keys to (session, exit_stack, loop) tuples
    self._sessions: Dict[
        str, tuple[ClientSession, AsyncExitStack, asyncio.AbstractEventLoop]
    ] = {}

    # Map of event loops to their respective locks to prevent race conditions
    # across different event loops in session creation.
    self._session_lock_map: dict[asyncio.AbstractEventLoop, asyncio.Lock] = {}
    self._lock_map_lock = threading.Lock()

  @property
  def _session_lock(self) -> asyncio.Lock:
    """Returns an asyncio.Lock bound to the current event loop."""
    current_loop = asyncio.get_running_loop()
    with self._lock_map_lock:
      if current_loop not in self._session_lock_map:
        self._session_lock_map[current_loop] = asyncio.Lock()
      return self._session_lock_map[current_loop]

  def _generate_session_key(
      self, merged_headers: Optional[Dict[str, str]] = None
  ) -> str:
    """Generates a session key based on connection params and merged headers.

    For StdioConnectionParams, returns a constant key since headers are not
    supported. For SSE and StreamableHTTP connections, generates a key based
    on the provided merged headers.

    Args:
        merged_headers: Already merged headers (base + additional).

    Returns:
        A unique session key string.
    """
    if isinstance(self._connection_params, StdioConnectionParams):
      # For stdio connections, headers are not supported, so use constant key
      return 'stdio_session'

    # For SSE and StreamableHTTP connections, use merged headers
    if merged_headers:
      headers_json = json.dumps(merged_headers, sort_keys=True)
      headers_hash = hashlib.md5(headers_json.encode()).hexdigest()
      return f'session_{headers_hash}'
    else:
      return 'session_no_headers'

  def _merge_headers(
      self, additional_headers: Optional[Dict[str, str]] = None
  ) -> Optional[Dict[str, str]]:
    """Merges base connection headers with additional headers.

    Args:
        additional_headers: Optional headers to merge with connection headers.

    Returns:
        Merged headers dictionary, or None if no headers are provided.
    """
    if isinstance(self._connection_params, StdioConnectionParams) or isinstance(
        self._connection_params, StdioServerParameters
    ):
      # Stdio connections don't support headers
      return None

    base_headers = {}
    if (
        hasattr(self._connection_params, 'headers')
        and self._connection_params.headers
    ):
      base_headers = self._connection_params.headers.copy()

    if additional_headers:
      base_headers.update(additional_headers)

    return base_headers

  def _is_session_disconnected(self, session: ClientSession) -> bool:
    """Checks if a session is disconnected or closed.

    Args:
        session: The ClientSession to check.

    Returns:
        True if the session is disconnected, False otherwise.
    """
    return session._read_stream._closed or session._write_stream._closed

  async def _cleanup_session(
      self,
      session_key: str,
      exit_stack: AsyncExitStack,
      stored_loop: asyncio.AbstractEventLoop,
  ):
    """Cleans up a session, handling different event loops safely.

    Args:
        session_key: The session key to clean up.
        exit_stack: The AsyncExitStack managing the session resources.
        stored_loop: The event loop on which the session was created.
    """
    current_loop = asyncio.get_running_loop()
    try:
      if stored_loop is current_loop:
        await exit_stack.aclose()
      elif stored_loop.is_closed():
        logger.warning(
            f'Error cleaning up session {session_key}: original event loop'
            ' is closed, resources may be leaked.'
        )
      else:
        # The old loop is still running in another thread;
        # schedule cleanup on it.
        logger.info(
            f'Scheduling cleanup of session {session_key} on its original'
            ' event loop.'
        )
        future = asyncio.run_coroutine_threadsafe(
            exit_stack.aclose(), stored_loop
        )

        # Attach a callback so errors don't go unnoticed
        def cleanup_done(f: asyncio.Future):
          try:
            if f.exception():
              logger.warning(
                  f'Error cleaning up session {session_key} on original'
                  f' loop: {f.exception()}'
              )
          except Exception as e:
            logger.warning(
                f'Failed to check cleanup status for {session_key}: {e}'
            )

        future.add_done_callback(cleanup_done)
    except Exception as e:
      logger.warning(
          f'Error during session cleanup for {session_key}: {e}',
          exc_info=True,
      )
    finally:
      if session_key in self._sessions:
        del self._sessions[session_key]

  def _create_client(self, merged_headers: Optional[Dict[str, str]] = None):
    """Creates an MCP client based on the connection parameters.

    Args:
        merged_headers: Optional headers to include in the connection.
                       Only applicable for SSE and StreamableHTTP connections.

    Returns:
        The appropriate MCP client instance.

    Raises:
        ValueError: If the connection parameters are not supported.
    """
    if isinstance(self._connection_params, StdioConnectionParams):
      client = stdio_client(
          server=self._connection_params.server_params,
          errlog=self._errlog,
      )
    elif isinstance(self._connection_params, SseConnectionParams):
      client = sse_client(
          url=self._connection_params.url,
          headers=merged_headers,
          timeout=self._connection_params.timeout,
          sse_read_timeout=self._connection_params.sse_read_timeout,
      )
    elif isinstance(self._connection_params, StreamableHTTPConnectionParams):
      client = streamablehttp_client(
          url=self._connection_params.url,
          headers=merged_headers,
          timeout=timedelta(seconds=self._connection_params.timeout),
          sse_read_timeout=timedelta(
              seconds=self._connection_params.sse_read_timeout
          ),
          terminate_on_close=self._connection_params.terminate_on_close,
          httpx_client_factory=self._connection_params.httpx_client_factory,
      )
    else:
      raise ValueError(
          'Unable to initialize connection. Connection should be'
          ' StdioServerParameters or SseServerParams, but got'
          f' {self._connection_params}'
      )
    return client

  async def create_session(
      self, headers: Optional[Dict[str, str]] = None
  ) -> ClientSession:
    """Creates and initializes an MCP client session.

    This method will check if an existing session for the given headers
    is still connected. If it's disconnected, it will be cleaned up and
    a new session will be created.

    Args:
        headers: Optional headers to include in the session. These will be
                merged with any existing connection headers. Only applicable
                for SSE and StreamableHTTP connections.

    Returns:
        ClientSession: The initialized MCP client session.
    """
    # Merge headers once at the beginning
    merged_headers = self._merge_headers(headers)

    # Generate session key using merged headers
    session_key = self._generate_session_key(merged_headers)

    # Use async lock to prevent race conditions
    async with self._session_lock:
      # Check if we have an existing session
      if session_key in self._sessions:
        session, exit_stack, stored_loop = self._sessions[session_key]

        # Check if the existing session is still connected and bound to the current loop
        current_loop = asyncio.get_running_loop()
        if stored_loop is current_loop and not self._is_session_disconnected(
            session
        ):
          # Session is still good, return it
          return session
        else:
          # Session is disconnected or from a different loop, clean it up
          logger.info(
              'Cleaning up session (disconnected or different loop): %s',
              session_key,
          )
          await self._cleanup_session(session_key, exit_stack, stored_loop)

      # Create a new session (either first time or replacing disconnected one)
      exit_stack = AsyncExitStack()
      timeout_in_seconds = (
          self._connection_params.timeout
          if hasattr(self._connection_params, 'timeout')
          else None
      )
      sse_read_timeout_in_seconds = (
          self._connection_params.sse_read_timeout
          if hasattr(self._connection_params, 'sse_read_timeout')
          else None
      )

      try:
        client = self._create_client(merged_headers)
        is_stdio = isinstance(self._connection_params, StdioConnectionParams)

        session = await asyncio.wait_for(
            exit_stack.enter_async_context(
                SessionContext(
                    client=client,
                    timeout=timeout_in_seconds,
                    sse_read_timeout=sse_read_timeout_in_seconds,
                    is_stdio=is_stdio,
                )
            ),
            timeout=timeout_in_seconds,
        )

        # Store session, exit stack, and loop in the pool
        self._sessions[session_key] = (
            session,
            exit_stack,
            asyncio.get_running_loop(),
        )
        logger.debug('Created new session: %s', session_key)
        return session

      except Exception as e:
        # If session creation fails, clean up the exit stack
        if exit_stack:
          try:
            await exit_stack.aclose()
          except Exception as exit_stack_error:
            logger.warning(
                'Error during session creation cleanup: %s', exit_stack_error
            )
        raise ConnectionError(f'Failed to create MCP session: {e}') from e

  def __getstate__(self):
    """Custom pickling to exclude non-picklable runtime objects."""
    state = self.__dict__.copy()
    # Remove unpicklable entries or those that shouldn't persist across pickle
    state['_sessions'] = {}
    state['_session_lock_map'] = {}

    # Locks and file-like objects cannot be pickled
    state.pop('_lock_map_lock', None)
    state.pop('_errlog', None)

    return state

  def __setstate__(self, state):
    """Custom unpickling to restore state."""
    self.__dict__.update(state)
    # Re-initialize members that were not pickled
    self._sessions = {}
    self._session_lock_map = {}
    self._lock_map_lock = threading.Lock()
    # If _errlog was removed during pickling, default to sys.stderr
    if not hasattr(self, '_errlog') or self._errlog is None:
      self._errlog = sys.stderr

  async def close(self):
    """Closes all sessions and cleans up resources."""
    async with self._session_lock:
      for session_key in list(self._sessions.keys()):
        _, exit_stack, stored_loop = self._sessions[session_key]
        await self._cleanup_session(session_key, exit_stack, stored_loop)


SseServerParams = SseConnectionParams

StreamableHTTPServerParams = StreamableHTTPConnectionParams
