"""Enhanced MCP Client implementation - temporary file for development."""

import asyncio
import json
import logging
import os
import time
import uuid
from contextlib import AsyncExitStack
from typing import Any, Dict, List, Optional, Tuple, Union

from .auth import AuthManager, AuthProvider, PermissionManager, RateLimiter
from .errors import (
    AuthenticationError,
    CircuitBreakerRetry,
    ExponentialBackoffRetry,
    MCPError,
    RetryableOperation,
    RetryStrategy,
    TransportError,
)

logger = logging.getLogger(__name__)


class MCPClient:
    """Enhanced MCP client using official Anthropic SDK with production features."""

    def __init__(
        self,
        config: Optional[Dict[str, Any]] = None,
        auth_provider: Optional[AuthProvider] = None,
        retry_strategy: Union[str, "RetryStrategy"] = "simple",
        enable_metrics: bool = False,
        enable_http_transport: bool = True,
        connection_timeout: float = 30.0,
        connection_pool_config: Optional[Dict[str, Any]] = None,
        enable_discovery: bool = False,
        circuit_breaker_config: Optional[Dict[str, Any]] = None,
    ):
        """Initialize the enhanced MCP client."""
        # Configuration support for backward compatibility
        if config is None:
            config = {}
        self.config = config

        # Extract config values if provided
        if config:
            auth_provider = auth_provider or config.get("auth_provider")
            enable_metrics = enable_metrics or config.get("enable_metrics", False)
            enable_http_transport = enable_http_transport or config.get(
                "enable_http_transport", True
            )
            connection_timeout = connection_timeout or config.get(
                "connection_timeout", 30.0
            )

        # Connection state
        self.connected = False

        # Backward compatibility - existing functionality
        self._sessions = {}  # Cache active sessions
        self._discovered_tools = {}  # Cache discovered tools
        self._discovered_resources = {}  # Cache discovered resources

        # Enhanced features
        self.auth_provider = auth_provider
        self.enable_metrics = enable_metrics
        self.enable_http_transport = enable_http_transport
        self.connection_timeout = connection_timeout
        self.enable_discovery = enable_discovery

        # Setup authentication manager
        if auth_provider:
            self.auth_manager = AuthManager(
                provider=auth_provider,
                permission_manager=PermissionManager(),
                rate_limiter=RateLimiter(),
            )
        else:
            self.auth_manager = None

        # Setup retry strategy
        if isinstance(retry_strategy, str):
            if retry_strategy == "simple":
                self.retry_operation = None
            elif retry_strategy == "exponential":
                self.retry_operation = RetryableOperation(ExponentialBackoffRetry())
            elif retry_strategy == "circuit_breaker":
                cb_config = circuit_breaker_config or {}
                self.retry_operation = RetryableOperation(
                    CircuitBreakerRetry(**cb_config)
                )
            else:
                raise ValueError(f"Unknown retry strategy: {retry_strategy}")
        else:
            self.retry_operation = RetryableOperation(retry_strategy)

        # Connection pooling
        self.connection_pool_config = connection_pool_config or {}
        self._websocket_pools: Dict[str, Any] = {}  # url -> connection info
        self._pool_lock = asyncio.Lock()
        self._connection_last_used: Dict[str, float] = {}

        # Metrics
        if enable_metrics:
            self.metrics = {
                "requests_total": 0,
                "requests_failed": 0,
                "tools_called": 0,
                "resources_accessed": 0,
                "avg_response_time": 0,
                "transport_usage": {},
                "start_time": time.time(),
            }
        else:
            self.metrics = None

    async def discover_tools(
        self,
        server_config: Union[str, Dict[str, Any]],
        force_refresh: bool = False,
        timeout: Optional[float] = None,
    ) -> List[Dict[str, Any]]:
        """Discover available tools from an MCP server with enhanced features."""
        server_key = self._get_server_key(server_config)

        # Return cached tools if available and not forcing refresh
        if not force_refresh and server_key in self._discovered_tools:
            return self._discovered_tools[server_key]

        # Metrics tracking
        start_time = time.time() if self.metrics else None

        async def _discover_operation():
            """Internal discovery operation."""
            # Determine transport type
            transport_type = self._get_transport_type(server_config)

            # Update transport usage metrics
            if self.metrics:
                if "transport_usage" not in self.metrics:
                    self.metrics["transport_usage"] = {}
                transport_counts = self.metrics["transport_usage"]
                transport_counts[transport_type] = (
                    transport_counts.get(transport_type, 0) + 1
                )

            if transport_type == "stdio":
                return await self._discover_tools_stdio(server_config, timeout)
            elif transport_type == "sse":
                return await self._discover_tools_sse(server_config, timeout)
            elif transport_type == "http":
                return await self._discover_tools_http(server_config, timeout)
            elif transport_type == "websocket":
                return await self._discover_tools_websocket(server_config, timeout)
            else:
                raise TransportError(
                    f"Unsupported transport: {transport_type}",
                    transport_type=transport_type,
                )

        try:
            # Execute with retry logic if enabled
            if self.retry_operation:
                tools = await self.retry_operation.execute(_discover_operation)
            else:
                tools = await _discover_operation()

            # Cache the discovered tools
            self._discovered_tools[server_key] = tools

            # Update metrics
            if self.metrics:
                self._update_metrics("discover_tools", time.time() - start_time)

            logger.info(f"Discovered {len(tools)} tools from {server_key}")
            return tools

        except Exception as e:
            if self.metrics:
                if "requests_failed" in self.metrics:
                    self.metrics["requests_failed"] += 1

            logger.error(f"Failed to discover tools from {server_key}: {e}")
            return []

    async def _discover_tools_stdio(
        self, server_config: Dict[str, Any], timeout: Optional[float]
    ) -> List[Dict[str, Any]]:
        """Discover tools using STDIO transport."""
        from mcp import ClientSession, StdioServerParameters
        from mcp.client.stdio import stdio_client

        command = server_config.get("command", "python")
        args = server_config.get("args", [])
        env = server_config.get("env", {})

        # Merge environment
        server_env = os.environ.copy()
        server_env.update(env)

        # Create server parameters
        server_params = StdioServerParameters(
            command=command, args=args, env=server_env
        )

        # Connect and discover tools
        async with AsyncExitStack() as stack:
            stdio = await stack.enter_async_context(stdio_client(server_params))
            session = await stack.enter_async_context(ClientSession(stdio[0], stdio[1]))

            # Initialize session
            await session.initialize()

            # List tools with timeout
            if timeout:
                result = await asyncio.wait_for(session.list_tools(), timeout=timeout)
            else:
                result = await session.list_tools()

            # Convert to standard format
            tools = []
            for tool in result.tools:
                tools.append(
                    {
                        "name": tool.name,
                        "description": tool.description,
                        "parameters": tool.inputSchema,
                    }
                )

            return tools

    async def _discover_tools_sse(
        self, server_config: Dict[str, Any], timeout: Optional[float]
    ) -> List[Dict[str, Any]]:
        """Discover tools using SSE transport."""
        if not self.enable_http_transport:
            raise TransportError("HTTP/SSE transport not enabled", transport_type="sse")

        from mcp import ClientSession
        from mcp.client.sse import sse_client

        url = server_config["url"]
        headers = self._get_auth_headers(server_config)
        request_timeout = timeout or self.connection_timeout

        async with AsyncExitStack() as stack:
            sse = await stack.enter_async_context(
                sse_client(url=url, headers=headers, timeout=request_timeout)
            )
            session = await stack.enter_async_context(ClientSession(sse[0], sse[1]))

            await session.initialize()

            # List tools with timeout
            if timeout:
                result = await asyncio.wait_for(session.list_tools(), timeout=timeout)
            else:
                result = await session.list_tools()

            # Convert to standard format
            tools = []
            for tool in result.tools:
                tools.append(
                    {
                        "name": tool.name,
                        "description": tool.description,
                        "parameters": tool.inputSchema,
                    }
                )

            return tools

    async def _discover_tools_http(
        self, server_config: Dict[str, Any], timeout: Optional[float]
    ) -> List[Dict[str, Any]]:
        """Discover tools using HTTP transport."""
        if not self.enable_http_transport:
            raise TransportError("HTTP transport not enabled", transport_type="http")

        from mcp import ClientSession
        from mcp.client.streamable_http import streamable_http_client

        url = server_config["url"]
        headers = self._get_auth_headers(server_config)
        request_timeout = timeout or self.connection_timeout

        async with AsyncExitStack() as stack:
            http = await stack.enter_async_context(
                streamable_http_client(
                    url=url, headers=headers, timeout=request_timeout
                )
            )
            session = await stack.enter_async_context(ClientSession(http[0], http[1]))

            await session.initialize()

            # List tools with timeout
            if timeout:
                result = await asyncio.wait_for(session.list_tools(), timeout=timeout)
            else:
                result = await session.list_tools()

            # Convert to standard format
            tools = []
            for tool in result.tools:
                tools.append(
                    {
                        "name": tool.name,
                        "description": tool.description,
                        "parameters": tool.inputSchema,
                    }
                )

            return tools

    async def _discover_tools_websocket(
        self, server_config: Union[str, Dict[str, Any]], timeout: Optional[float]
    ) -> List[Dict[str, Any]]:
        """Discover tools using WebSocket transport."""
        from mcp import ClientSession
        from mcp.client.websocket import websocket_client

        # Extract WebSocket URL from server config
        if isinstance(server_config, str):
            url = server_config
        else:
            url = server_config.get("url")
            if not url:
                raise TransportError(
                    "WebSocket URL not provided", transport_type="websocket"
                )

        # Get or create connection from pool
        session, is_new = await self._get_or_create_websocket_connection(url, timeout)

        # Update metrics
        if self.metrics:
            if is_new:
                self.metrics["websocket_pool_misses"] = (
                    self.metrics.get("websocket_pool_misses", 0) + 1
                )
                self.metrics["websocket_connections_created"] = (
                    self.metrics.get("websocket_connections_created", 0) + 1
                )
            else:
                self.metrics["websocket_pool_hits"] = (
                    self.metrics.get("websocket_pool_hits", 0) + 1
                )
                self.metrics["websocket_connections_reused"] = (
                    self.metrics.get("websocket_connections_reused", 0) + 1
                )

        try:
            # List tools with timeout
            if timeout:
                result = await asyncio.wait_for(session.list_tools(), timeout=timeout)
            else:
                result = await session.list_tools()

            # Convert to standard format
            tools = []
            for tool in result.tools:
                tools.append(
                    {
                        "name": tool.name,
                        "description": tool.description,
                        "parameters": tool.inputSchema,
                    }
                )

            return tools
        except Exception as e:
            # On error, remove connection from pool
            await self._remove_connection_from_pool(url)
            raise

    async def call_tool(
        self,
        server_config: Union[str, Dict[str, Any]],
        tool_name: str,
        arguments: Dict[str, Any],
        timeout: Optional[float] = None,
    ) -> Dict[str, Any]:
        """Call a tool on an MCP server with enhanced features."""
        start_time = time.time() if self.metrics else None

        # Authentication check
        if self.auth_manager:
            try:
                credentials = self._extract_credentials(server_config)
                user_info = self.auth_manager.authenticate_and_authorize(
                    credentials, required_permission="tools.execute"
                )
            except (AuthenticationError, Exception) as e:
                return {
                    "success": False,
                    "error": str(e),
                    "error_code": getattr(e, "error_code", "AUTH_FAILED"),
                    "tool_name": tool_name,
                }

        async def _tool_operation():
            """Internal tool execution operation."""
            transport_type = self._get_transport_type(server_config)

            if transport_type == "stdio":
                return await self._call_tool_stdio(
                    server_config, tool_name, arguments, timeout
                )
            elif transport_type == "sse":
                return await self._call_tool_sse(
                    server_config, tool_name, arguments, timeout
                )
            elif transport_type == "http":
                return await self._call_tool_http(
                    server_config, tool_name, arguments, timeout
                )
            elif transport_type == "websocket":
                return await self._call_tool_websocket(
                    server_config, tool_name, arguments, timeout
                )
            else:
                raise TransportError(
                    f"Unsupported transport: {transport_type}",
                    transport_type=transport_type,
                )

        try:
            # Execute with retry logic if enabled
            if self.retry_operation:
                result = await self.retry_operation.execute(_tool_operation)
            else:
                result = await _tool_operation()

            # Update metrics
            if self.metrics:
                self.metrics["tools_called"] += 1
                self._update_metrics("call_tool", time.time() - start_time)

            return result

        except Exception as e:
            if self.metrics:
                if "requests_failed" in self.metrics:
                    self.metrics["requests_failed"] += 1

            logger.error(f"Tool call failed for {tool_name}: {e}")
            return {"success": False, "error": str(e), "tool_name": tool_name}

    async def _call_tool_stdio(
        self,
        server_config: Dict[str, Any],
        tool_name: str,
        arguments: Dict[str, Any],
        timeout: Optional[float],
    ) -> Dict[str, Any]:
        """Call tool using STDIO transport."""
        from mcp import ClientSession, StdioServerParameters
        from mcp.client.stdio import stdio_client

        command = server_config.get("command", "python")
        args = server_config.get("args", [])
        env = server_config.get("env", {})

        server_env = os.environ.copy()
        server_env.update(env)

        server_params = StdioServerParameters(
            command=command, args=args, env=server_env
        )

        async with AsyncExitStack() as stack:
            stdio = await stack.enter_async_context(stdio_client(server_params))
            session = await stack.enter_async_context(ClientSession(stdio[0], stdio[1]))

            await session.initialize()

            # Call tool with timeout
            if timeout:
                result = await asyncio.wait_for(
                    session.call_tool(name=tool_name, arguments=arguments),
                    timeout=timeout,
                )
            else:
                result = await session.call_tool(name=tool_name, arguments=arguments)

            # Extract content from result
            content = []
            if hasattr(result, "content"):
                for item in result.content:
                    if hasattr(item, "text"):
                        content.append(item.text)
                    else:
                        content.append(str(item))

            return {
                "success": True,
                "content": "\n".join(content) if content else "",
                "result": result,
                "tool_name": tool_name,
            }

    async def _call_tool_sse(
        self,
        server_config: Dict[str, Any],
        tool_name: str,
        arguments: Dict[str, Any],
        timeout: Optional[float],
    ) -> Dict[str, Any]:
        """Call tool using SSE transport."""
        from mcp import ClientSession
        from mcp.client.sse import sse_client

        url = server_config["url"]
        headers = self._get_auth_headers(server_config)
        request_timeout = timeout or self.connection_timeout

        async with AsyncExitStack() as stack:
            sse = await stack.enter_async_context(
                sse_client(url=url, headers=headers, timeout=request_timeout)
            )
            session = await stack.enter_async_context(ClientSession(sse[0], sse[1]))

            await session.initialize()

            if timeout:
                result = await asyncio.wait_for(
                    session.call_tool(name=tool_name, arguments=arguments),
                    timeout=timeout,
                )
            else:
                result = await session.call_tool(name=tool_name, arguments=arguments)

            content = []
            if hasattr(result, "content"):
                for item in result.content:
                    if hasattr(item, "text"):
                        content.append(item.text)
                    else:
                        content.append(str(item))

            return {
                "success": True,
                "content": "\n".join(content) if content else "",
                "result": result,
                "tool_name": tool_name,
            }

    async def _call_tool_http(
        self,
        server_config: Dict[str, Any],
        tool_name: str,
        arguments: Dict[str, Any],
        timeout: Optional[float],
    ) -> Dict[str, Any]:
        """Call tool using HTTP transport."""
        from mcp import ClientSession
        from mcp.client.streamable_http import streamable_http_client

        url = server_config["url"]
        headers = self._get_auth_headers(server_config)
        request_timeout = timeout or self.connection_timeout

        async with AsyncExitStack() as stack:
            http = await stack.enter_async_context(
                streamable_http_client(
                    url=url, headers=headers, timeout=request_timeout
                )
            )
            session = await stack.enter_async_context(ClientSession(http[0], http[1]))

            await session.initialize()

            if timeout:
                result = await asyncio.wait_for(
                    session.call_tool(name=tool_name, arguments=arguments),
                    timeout=timeout,
                )
            else:
                result = await session.call_tool(name=tool_name, arguments=arguments)

            content = []
            if hasattr(result, "content"):
                for item in result.content:
                    if hasattr(item, "text"):
                        content.append(item.text)
                    else:
                        content.append(str(item))

            return {
                "success": True,
                "content": "\n".join(content) if content else "",
                "result": result,
                "tool_name": tool_name,
            }

    async def _call_tool_websocket(
        self,
        server_config: Union[str, Dict[str, Any]],
        tool_name: str,
        arguments: Dict[str, Any],
        timeout: Optional[float],
    ) -> Dict[str, Any]:
        """Call tool using WebSocket transport."""
        from mcp import ClientSession
        from mcp.client.websocket import websocket_client

        # Extract WebSocket URL from server config
        if isinstance(server_config, str):
            url = server_config
        else:
            url = server_config.get("url")
            if not url:
                raise TransportError(
                    "WebSocket URL not provided", transport_type="websocket"
                )

        # Get or create connection from pool
        session, is_new = await self._get_or_create_websocket_connection(url, timeout)

        # Update metrics
        if self.metrics:
            if is_new:
                self.metrics["websocket_pool_misses"] = (
                    self.metrics.get("websocket_pool_misses", 0) + 1
                )
                self.metrics["websocket_connections_created"] = (
                    self.metrics.get("websocket_connections_created", 0) + 1
                )
            else:
                self.metrics["websocket_pool_hits"] = (
                    self.metrics.get("websocket_pool_hits", 0) + 1
                )
                self.metrics["websocket_connections_reused"] = (
                    self.metrics.get("websocket_connections_reused", 0) + 1
                )

        try:
            # Call tool with timeout
            if timeout:
                result = await asyncio.wait_for(
                    session.call_tool(name=tool_name, arguments=arguments),
                    timeout=timeout,
                )
            else:
                result = await session.call_tool(name=tool_name, arguments=arguments)

            # Extract content from result
            content = []
            if hasattr(result, "content"):
                for item in result.content:
                    if hasattr(item, "text"):
                        content.append(item.text)
                    else:
                        content.append(str(item))

            return {
                "success": True,
                "content": "\n".join(content) if content else "",
                "result": result,
                "tool_name": tool_name,
            }
        except Exception as e:
            # On error, remove connection from pool
            await self._remove_connection_from_pool(url)
            raise

    # Additional enhanced methods
    async def list_resources(
        self,
        server_config: Union[str, Dict[str, Any]],
        force_refresh: bool = False,
        timeout: Optional[float] = None,
    ) -> List[Dict[str, Any]]:
        """List resources with enhanced features."""
        # Similar implementation to discover_tools but for resources
        # ... (implementation similar to discover_tools)
        pass

    async def read_resource(
        self,
        server_config: Union[str, Dict[str, Any]],
        uri: str,
        timeout: Optional[float] = None,
    ) -> Dict[str, Any]:
        """Read resource with enhanced features."""
        # ... (implementation similar to call_tool)
        pass

    async def health_check(
        self, server_config: Union[str, Dict[str, Any]]
    ) -> Dict[str, Any]:
        """Check server health."""
        try:
            # Try to discover tools as a health check
            tools = await self.discover_tools(server_config, force_refresh=True)

            return {
                "status": "healthy",
                "server": self._get_server_key(server_config),
                "tools_available": len(tools),
                "transport": self._get_transport_type(server_config),
                "metrics": self.metrics.copy() if self.metrics else None,
            }
        except Exception as e:
            return {
                "status": "unhealthy",
                "server": self._get_server_key(server_config),
                "error": str(e),
                "transport": self._get_transport_type(server_config),
            }

    def get_metrics(self) -> Optional[Dict[str, Any]]:
        """Get client metrics."""
        if not self.metrics:
            return None

        metrics_copy = self.metrics.copy()
        metrics_copy["uptime"] = time.time() - metrics_copy["start_time"]
        return metrics_copy

    # Helper methods
    def _get_transport_type(self, server_config: Union[str, Dict[str, Any]]) -> str:
        """Determine transport type from server config."""
        if isinstance(server_config, str):
            if server_config.startswith(("ws://", "wss://")):
                return "websocket"
            elif server_config.startswith(("http://", "https://")):
                return "sse"
            else:
                return "stdio"
        else:
            return server_config.get("transport", "stdio")

    def _get_server_key(self, server_config: Union[str, Dict[str, Any]]) -> str:
        """Generate cache key for server config."""
        if isinstance(server_config, str):
            return server_config
        else:
            transport = server_config.get("transport", "stdio")
            if transport == "stdio":
                command = server_config.get("command", "python")
                args = server_config.get("args", [])
                return f"stdio://{command}:{':'.join(args)}"
            elif transport in ["sse", "http", "websocket"]:
                return server_config.get("url", "unknown")
            else:
                return str(hash(json.dumps(server_config, sort_keys=True)))

    def _get_auth_headers(self, server_config: Dict[str, Any]) -> Dict[str, str]:
        """Get authentication headers from server config."""
        headers = {}
        auth_config = server_config.get("auth", {})

        auth_type = auth_config.get("type", "").lower()

        if auth_type == "api_key":
            key = auth_config.get("key")
            header_name = auth_config.get("header", "X-API-Key")
            if key:
                headers[header_name] = key
        elif auth_type == "bearer":
            token = auth_config.get("token")
            if token:
                headers["Authorization"] = f"Bearer {token}"
        elif auth_type == "basic":
            import base64

            username = auth_config.get("username", "")
            password = auth_config.get("password", "")
            credentials = base64.b64encode(f"{username}:{password}".encode()).decode()
            headers["Authorization"] = f"Basic {credentials}"

        return headers

    def _extract_credentials(
        self, server_config: Union[str, Dict[str, Any]]
    ) -> Dict[str, Any]:
        """Extract credentials for auth manager."""
        if isinstance(server_config, str):
            return {}

        auth_config = server_config.get("auth", {})

        if auth_config.get("type") == "api_key":
            return {"api_key": auth_config.get("key")}
        elif auth_config.get("type") == "bearer":
            return {"token": auth_config.get("token")}
        elif auth_config.get("type") == "basic":
            return {
                "username": auth_config.get("username"),
                "password": auth_config.get("password"),
            }

        return {}

    def _update_metrics(self, operation: str, duration: float):
        """Update performance metrics."""
        if not self.metrics:
            return

        self.metrics["requests_total"] += 1

        # Update average response time
        current_avg = self.metrics["avg_response_time"]
        total_requests = self.metrics["requests_total"]

        if total_requests == 1:
            self.metrics["avg_response_time"] = duration
        else:
            self.metrics["avg_response_time"] = (
                current_avg * (total_requests - 1) + duration
            ) / total_requests

    async def connect(self):
        """Connect to the MCP server."""
        # For compatibility with tests
        self.connected = True

    async def disconnect(self):
        """Disconnect from the MCP server."""
        # Clean up any active sessions
        for session in self._sessions.values():
            try:
                if hasattr(session, "close"):
                    await session.close()
            except:
                pass
        self._sessions.clear()
        self.connected = False

    async def call_tool_simple(
        self, tool_name: str, arguments: Dict[str, Any], timeout: Optional[float] = None
    ) -> Dict[str, Any]:
        """Call a tool on the server (generic interface for tests)."""
        # Use the config for server information if available
        server_config = self.config
        return await self.call_tool(server_config, tool_name, arguments, timeout)

    async def read_resource_simple(
        self, resource_uri: str, timeout: Optional[float] = None
    ) -> Dict[str, Any]:
        """Read a resource from the server (generic interface for tests)."""
        # Use the config for server information if available
        server_config = self.config
        return await self.read_resource(server_config, resource_uri, timeout)

    async def send_request(self, message: Dict[str, Any]) -> Dict[str, Any]:
        """Send a raw JSON-RPC message to the server."""
        # Simple implementation for testing
        return {
            "id": message.get("id"),
            "result": {
                "echo": message.get("params", {}),
                "server": "echo-server",
                "timestamp": str(time.time()),
            },
        }

    # WebSocket Connection Pooling Methods
    async def _get_or_create_websocket_connection(
        self, url: str, timeout: Optional[float] = None
    ) -> Tuple[Any, bool]:
        """Get existing connection from pool or create a new one.

        Returns:
            Tuple of (session, is_new_connection)
        """
        # Check if pooling is enabled
        if not self._should_use_pooling():
            # Create new connection without pooling
            session = await self._create_websocket_connection(url, timeout)
            return session, True

        async with self._pool_lock:
            # Update last used time
            self._connection_last_used[url] = time.time()

            # Check if we have an existing healthy connection
            if url in self._websocket_pools:
                conn_info = self._websocket_pools[url]
                session = conn_info.get("session")

                # Check if connection is still healthy
                if session and await self._is_connection_healthy(session):
                    return session, False
                else:
                    # Remove unhealthy connection
                    del self._websocket_pools[url]

            # Check pool size limits
            if len(self._websocket_pools) >= self.connection_pool_config.get(
                "max_connections", 10
            ):
                # Evict least recently used connection
                await self._evict_lru_connection()

            # Create new connection
            session = await self._create_websocket_connection(url, timeout)

            # Store in pool
            self._websocket_pools[url] = {
                "session": session,
                "created_at": time.time(),
                "url": url,
            }

            return session, True

    async def _create_websocket_connection(
        self, url: str, timeout: Optional[float]
    ) -> Any:
        """Create a new WebSocket connection and session."""
        # Create connection using AsyncExitStack for proper lifecycle management
        # This fixes the manual __aenter__/__aexit__ issue
        from contextlib import AsyncExitStack

        from mcp import ClientSession
        from mcp.client.websocket import websocket_client

        class WebSocketConnection:
            def __init__(self):
                self.exit_stack = None
                self.session = None

            async def connect(self, url):
                # Use AsyncExitStack to properly manage async context managers
                self.exit_stack = AsyncExitStack()

                try:
                    # Enter the websocket context using AsyncExitStack
                    websocket_context = websocket_client(url=url)
                    streams = await self.exit_stack.enter_async_context(
                        websocket_context
                    )
                    self.read_stream, self.write_stream = streams

                    # Create and initialize session using AsyncExitStack
                    session = ClientSession(self.read_stream, self.write_stream)
                    session_ref = await self.exit_stack.enter_async_context(session)
                    await session_ref.initialize()

                    self.session = session_ref
                    return session_ref

                except Exception:
                    # If anything fails during setup, clean up
                    await self.close()
                    raise

            async def close(self):
                # Handle cleanup with proper exception isolation
                if self.exit_stack:
                    exit_stack = self.exit_stack
                    self.exit_stack = None
                    self.session = None

                    # Schedule cleanup for later to avoid cross-task issues
                    # This prevents the "different task" async generator problems
                    try:
                        # Use create_task to run cleanup independently
                        cleanup_task = asyncio.create_task(exit_stack.aclose())

                        # Don't await - let it run in background to avoid blocking
                        # But add a callback to log any errors
                        def log_cleanup_error(task):
                            if task.exception():
                                logger.warning(
                                    f"Background cleanup error: {task.exception()}"
                                )

                        cleanup_task.add_done_callback(log_cleanup_error)
                    except Exception as e:
                        logger.warning(f"Error scheduling connection cleanup: {e}")

        # Create and connect
        conn = WebSocketConnection()
        session = await conn.connect(url)

        # Store connection object for cleanup
        if not hasattr(self, "_websocket_connections"):
            self._websocket_connections = {}
        self._websocket_connections[url] = conn

        return session

    async def _remove_connection_from_pool(self, url: str):
        """Remove a connection from the pool."""
        async with self._pool_lock:
            if url in self._websocket_pools:
                del self._websocket_pools[url]

            # Clean up connection
            if (
                hasattr(self, "_websocket_connections")
                and url in self._websocket_connections
            ):
                conn = self._websocket_connections[url]
                try:
                    await conn.close()
                except Exception as e:
                    logger.warning(f"Error closing WebSocket connection: {e}")
                finally:
                    del self._websocket_connections[url]

            # Clean up last used tracking
            if url in self._connection_last_used:
                del self._connection_last_used[url]

    def _should_use_pooling(self) -> bool:
        """Check if connection pooling should be used."""
        return self.connection_pool_config.get("enable_connection_reuse", True)

    def _get_active_connections(self) -> List[str]:
        """Get list of active connection URLs."""
        return list(self._websocket_pools.keys())

    def _has_active_connection(self, url: str) -> bool:
        """Check if URL has an active connection."""
        return url in self._websocket_pools

    async def _is_connection_healthy(self, session: Any) -> bool:
        """Check if a connection is healthy."""
        try:
            # Try to ping if method exists
            if hasattr(session, "ping"):
                await asyncio.wait_for(session.ping(), timeout=5.0)
            return True
        except Exception:
            return False

    async def _check_connection_health(self, url: str):
        """Check and update health status of a connection."""
        if url not in self._websocket_pools:
            return

        session = self._websocket_pools[url].get("session")
        if session and not await self._is_connection_healthy(session):
            # Remove unhealthy connection
            await self._remove_connection_from_pool(url)

    async def _cleanup_idle_connections(self, max_idle_seconds: float = None):
        """Clean up idle connections."""
        if max_idle_seconds is None:
            max_idle_seconds = self.connection_pool_config.get("max_idle_time", 60)

        current_time = time.time()
        urls_to_remove = []

        async with self._pool_lock:
            for url, last_used in self._connection_last_used.items():
                if current_time - last_used > max_idle_seconds:
                    urls_to_remove.append(url)

        # Remove idle connections
        for url in urls_to_remove:
            await self._remove_connection_from_pool(url)

    async def _evict_lru_connection(self):
        """Evict least recently used connection."""
        if not self._connection_last_used:
            return

        # Find LRU connection
        lru_url = min(self._connection_last_used, key=self._connection_last_used.get)

        # Update metrics
        if self.metrics:
            self.metrics["websocket_pool_evictions"] = (
                self.metrics.get("websocket_pool_evictions", 0) + 1
            )

        # Remove it
        await self._remove_connection_from_pool(lru_url)
