"""
Complete MCP Transport Implementations.

This module provides comprehensive transport layer implementations for MCP,
including enhanced STDIO, SSE, StreamableHTTP, and WebSocket transports.
All implementations build on the official MCP Python SDK while adding
production-ready features like security, connection management, and monitoring.

Features:
- Enhanced STDIO transport with proper process management
- Complete SSE transport with endpoint negotiation
- StreamableHTTP transport with session management
- WebSocket transport for real-time communication
- Transport security and validation
- Connection pooling and management
- Health checking and monitoring

Examples:
    Enhanced STDIO transport:

    >>> from kailash.mcp_server.transports import EnhancedStdioTransport
    >>> transport = EnhancedStdioTransport(
    ...     command="python",
    ...     args=["-m", "my_mcp_server"],
    ...     environment_filter=["PATH", "PYTHONPATH"]
    ... )
    >>> async with transport:
    ...     session = await transport.create_session()

    SSE transport with security:

    >>> from kailash.mcp_server.transports import SSETransport
    >>> transport = SSETransport(
    ...     base_url="https://api.example.com/mcp",
    ...     auth_header="Bearer token123",
    ...     validate_origin=True
    ... )

    StreamableHTTP with session management:

    >>> from kailash.mcp_server.transports import StreamableHTTPTransport
    >>> transport = StreamableHTTPTransport(
    ...     base_url="https://api.example.com/mcp",
    ...     session_management=True,
    ...     streaming_threshold=1024
    ... )
"""

import asyncio
import json
import logging
import os
import platform
import signal
import socket
import subprocess
import time
import uuid
import weakref
from abc import ABC, abstractmethod
from contextlib import AsyncExitStack
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
from urllib.parse import urljoin, urlparse

import aiohttp
import websockets

from .auth import AuthProvider
from .errors import MCPError, MCPErrorCode, TransportError
from .protocol import MetaData, ProtocolManager

logger = logging.getLogger(__name__)


class TransportSecurity:
    """Security utilities for MCP transports."""

    ALLOWED_SCHEMES = {"http", "https", "ws", "wss"}
    BLOCKED_HOSTS = {
        "169.254.169.254",
        "localhost",
        "127.0.0.1",
    }  # Basic DNS rebinding protection

    @classmethod
    def validate_url(cls, url: str, allow_localhost: bool = False) -> bool:
        """Validate URL for security.

        Args:
            url: URL to validate
            allow_localhost: Allow localhost connections

        Returns:
            True if URL is safe
        """
        try:
            parsed = urlparse(url)

            # Check scheme
            if parsed.scheme not in cls.ALLOWED_SCHEMES:
                logger.warning(f"Blocked unsafe scheme: {parsed.scheme}")
                return False

            # Check host
            if not allow_localhost and parsed.hostname in cls.BLOCKED_HOSTS:
                logger.warning(f"Blocked potentially unsafe host: {parsed.hostname}")
                return False

            # Check for IP address patterns that could be exploited
            if parsed.hostname and parsed.hostname.startswith("0."):
                logger.warning(f"Blocked potentially unsafe IP: {parsed.hostname}")
                return False

            return True

        except Exception as e:
            logger.error(f"URL validation error: {e}")
            return False

    @classmethod
    def validate_origin(cls, origin: str, expected_origins: List[str]) -> bool:
        """Validate request origin.

        Args:
            origin: Request origin
            expected_origins: List of allowed origins

        Returns:
            True if origin is allowed
        """
        if not origin:
            return False

        # Exact match
        if origin in expected_origins:
            return True

        # Wildcard patterns
        for expected in expected_origins:
            if "*" in expected:
                # Convert wildcard pattern to regex
                import re

                pattern = expected.replace(".", r"\.").replace("*", ".*")
                if re.match(f"^{pattern}$", origin):
                    return True

        return False


class BaseTransport(ABC):
    """Base class for MCP transports."""

    def __init__(
        self,
        name: str,
        auth_provider: Optional[AuthProvider] = None,
        timeout: float = 30.0,
        max_retries: int = 3,
        enable_metrics: bool = True,
    ):
        """Initialize base transport.

        Args:
            name: Transport name
            auth_provider: Authentication provider
            timeout: Connection timeout
            max_retries: Maximum retry attempts
            enable_metrics: Enable metrics collection
        """
        self.name = name
        self.auth_provider = auth_provider
        self.timeout = timeout
        self.max_retries = max_retries
        self.enable_metrics = enable_metrics

        # State
        self._connected = False
        self._sessions: Set[weakref.ref] = set()
        self._metrics: Dict[str, Any] = {}

        # Initialize metrics
        if enable_metrics:
            self._metrics = {
                "connections_total": 0,
                "connections_failed": 0,
                "messages_sent": 0,
                "messages_received": 0,
                "bytes_sent": 0,
                "bytes_received": 0,
                "errors_total": 0,
                "start_time": time.time(),
            }

    @abstractmethod
    async def connect(self) -> None:
        """Connect the transport."""
        pass

    @abstractmethod
    async def disconnect(self) -> None:
        """Disconnect the transport."""
        pass

    @abstractmethod
    async def send_message(self, message: Dict[str, Any]) -> None:
        """Send a message."""
        pass

    @abstractmethod
    async def receive_message(self) -> Dict[str, Any]:
        """Receive a message."""
        pass

    async def __aenter__(self):
        """Async context manager entry."""
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit."""
        await self.disconnect()

    def is_connected(self) -> bool:
        """Check if transport is connected."""
        return self._connected

    def get_metrics(self) -> Dict[str, Any]:
        """Get transport metrics."""
        if not self.enable_metrics:
            return {}

        metrics = self._metrics.copy()
        metrics["uptime"] = time.time() - metrics["start_time"]
        metrics["active_sessions"] = len(self._sessions)
        return metrics

    def _update_metrics(self, metric: str, value: Union[int, float] = 1):
        """Update metrics."""
        if self.enable_metrics and metric in self._metrics:
            self._metrics[metric] += value


class EnhancedStdioTransport(BaseTransport):
    """Enhanced STDIO transport with proper process management."""

    def __init__(
        self,
        command: str,
        args: Optional[List[str]] = None,
        env: Optional[Dict[str, str]] = None,
        working_directory: Optional[str] = None,
        environment_filter: Optional[List[str]] = None,
        **kwargs,
    ):
        """Initialize enhanced STDIO transport.

        Args:
            command: Command to execute
            args: Command arguments
            env: Environment variables
            working_directory: Working directory
            environment_filter: Allowed environment variables
            **kwargs: Base transport arguments
        """
        super().__init__("stdio", **kwargs)

        self.command = command
        self.args = args or []
        self.env = env or {}
        self.working_directory = working_directory
        self.environment_filter = environment_filter

        # Process management
        self.process: Optional[asyncio.subprocess.Process] = None
        self._read_task: Optional[asyncio.Task] = None
        self._write_queue: asyncio.Queue = asyncio.Queue()
        self._message_buffer: List[str] = []

    async def connect(self) -> None:
        """Start the subprocess and connect."""
        if self._connected:
            return

        try:
            # Prepare environment
            process_env = self._prepare_environment()

            # Start process
            self.process = await asyncio.create_subprocess_exec(
                self.command,
                *self.args,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=process_env,
                cwd=self.working_directory,
            )

            # Start I/O tasks
            self._read_task = asyncio.create_task(self._read_loop())

            self._connected = True
            self._update_metrics("connections_total")

            logger.info(f"STDIO transport connected: {self.command}")

        except Exception as e:
            self._update_metrics("connections_failed")
            raise TransportError(
                f"Failed to start process: {e}", transport_type="stdio"
            )

    async def disconnect(self) -> None:
        """Terminate the subprocess."""
        if not self._connected:
            return

        self._connected = False

        # Cancel read task
        if self._read_task and not self._read_task.done():
            self._read_task.cancel()
            try:
                await self._read_task
            except asyncio.CancelledError:
                pass

        # Terminate process
        if self.process:
            try:
                # Try graceful termination first
                self.process.terminate()

                # Wait with timeout
                try:
                    await asyncio.wait_for(self.process.wait(), timeout=5.0)
                except asyncio.TimeoutError:
                    # Force kill if needed
                    if platform.system() != "Windows":
                        self.process.kill()
                    await self.process.wait()

            except Exception as e:
                logger.error(f"Error terminating process: {e}")

            finally:
                self.process = None

        logger.info("STDIO transport disconnected")

    async def send_message(self, message: Dict[str, Any]) -> None:
        """Send message to subprocess."""
        if not self._connected or not self.process:
            raise TransportError("Transport not connected", transport_type="stdio")

        try:
            # Serialize message
            message_data = json.dumps(message) + "\n"
            message_bytes = message_data.encode("utf-8")

            # Send to subprocess
            self.process.stdin.write(message_bytes)
            await self.process.stdin.drain()

            self._update_metrics("messages_sent")
            self._update_metrics("bytes_sent", len(message_bytes))

        except Exception as e:
            self._update_metrics("errors_total")
            raise TransportError(f"Failed to send message: {e}", transport_type="stdio")

    async def receive_message(self) -> Dict[str, Any]:
        """Receive message from subprocess."""
        if not self._connected:
            raise TransportError("Transport not connected", transport_type="stdio")

        try:
            # Wait for message in buffer
            while not self._message_buffer:
                await asyncio.sleep(0.01)  # Small delay to prevent busy waiting

                if not self._connected:
                    raise TransportError(
                        "Transport disconnected", transport_type="stdio"
                    )

            # Get message from buffer
            message_data = self._message_buffer.pop(0)
            message = json.loads(message_data)

            self._update_metrics("messages_received")
            self._update_metrics("bytes_received", len(message_data))

            return message

        except json.JSONDecodeError as e:
            self._update_metrics("errors_total")
            raise TransportError(f"Invalid JSON received: {e}", transport_type="stdio")
        except Exception as e:
            self._update_metrics("errors_total")
            raise TransportError(
                f"Failed to receive message: {e}", transport_type="stdio"
            )

    async def _read_loop(self):
        """Background task to read from subprocess."""
        if not self.process:
            return

        try:
            while self._connected and self.process:
                # Read line from stdout
                line = await self.process.stdout.readline()

                if not line:
                    # Process ended
                    break

                # Decode and strip
                line_str = line.decode("utf-8").strip()

                if line_str:
                    self._message_buffer.append(line_str)

        except Exception as e:
            logger.error(f"STDIO read loop error: {e}")
        finally:
            if self._connected:
                await self.disconnect()

    def _prepare_environment(self) -> Dict[str, str]:
        """Prepare process environment variables."""
        # Start with filtered parent environment
        if self.environment_filter:
            process_env = {
                key: value
                for key, value in os.environ.items()
                if key in self.environment_filter
            }
        else:
            process_env = os.environ.copy()

        # Add custom environment variables
        process_env.update(self.env)

        return process_env

    async def get_process_info(self) -> Dict[str, Any]:
        """Get information about the subprocess."""
        if not self.process:
            return {}

        return {
            "pid": self.process.pid,
            "returncode": self.process.returncode,
            "command": [self.command] + self.args,
            "working_directory": self.working_directory,
        }


class SSETransport(BaseTransport):
    """Server-Sent Events transport with endpoint negotiation."""

    def __init__(
        self,
        base_url: str,
        auth_header: Optional[str] = None,
        validate_origin: bool = True,
        allowed_origins: Optional[List[str]] = None,
        endpoint_path: str = "/sse",
        message_path: str = "/message",
        allow_localhost: bool = False,
        skip_security_validation: bool = False,
        **kwargs,
    ):
        """Initialize SSE transport.

        Args:
            base_url: Base URL for the server
            auth_header: Authorization header
            validate_origin: Enable origin validation
            allowed_origins: List of allowed origins
            endpoint_path: SSE endpoint path
            message_path: Message posting path
            allow_localhost: Allow connections to localhost (for testing)
            skip_security_validation: Skip all security validation (for testing)
            **kwargs: Base transport arguments
        """
        super().__init__("sse", **kwargs)

        self.base_url = base_url.rstrip("/")
        self.auth_header = auth_header
        self.validate_origin = validate_origin
        self.allowed_origins = allowed_origins or [base_url]
        self.endpoint_path = endpoint_path
        self.message_path = message_path
        self.allow_localhost = allow_localhost
        self.skip_security_validation = skip_security_validation

        # Connection state
        self.session: Optional[aiohttp.ClientSession] = None
        self.sse_response: Optional[aiohttp.ClientResponse] = None
        self._read_task: Optional[asyncio.Task] = None
        self._message_queue: asyncio.Queue = asyncio.Queue()

    async def connect(self) -> None:
        """Connect to SSE endpoint."""
        if self._connected:
            return

        # Validate URL (with configurable security)
        if not self.skip_security_validation:
            if not TransportSecurity.validate_url(
                self.base_url, allow_localhost=self.allow_localhost
            ):
                raise TransportError("Invalid or unsafe URL", transport_type="sse")

        try:
            # Create session
            headers = {}
            if self.auth_header:
                headers["Authorization"] = self.auth_header

            # Add CORS headers if origin validation is enabled
            if self.validate_origin:
                headers["Origin"] = self.base_url

            timeout = aiohttp.ClientTimeout(total=self.timeout)
            self.session = aiohttp.ClientSession(headers=headers, timeout=timeout)

            # Connect to SSE endpoint
            sse_url = urljoin(self.base_url, self.endpoint_path)
            self.sse_response = await self.session.get(
                sse_url,
                headers={"Accept": "text/event-stream", "Cache-Control": "no-cache"},
            )

            if self.sse_response.status != 200:
                raise TransportError(
                    f"SSE connection failed: {self.sse_response.status}",
                    transport_type="sse",
                )

            # Start reading SSE events
            self._read_task = asyncio.create_task(self._read_sse_events())

            self._connected = True
            self._update_metrics("connections_total")

            logger.info(f"SSE transport connected: {sse_url}")

        except Exception as e:
            self._update_metrics("connections_failed")
            await self._cleanup_connection()
            raise TransportError(f"SSE connection failed: {e}", transport_type="sse")

    async def disconnect(self) -> None:
        """Disconnect from SSE endpoint."""
        if not self._connected:
            return

        self._connected = False
        await self._cleanup_connection()

        logger.info("SSE transport disconnected")

    async def send_message(self, message: Dict[str, Any]) -> None:
        """Send message via HTTP POST."""
        if not self._connected or not self.session:
            raise TransportError("Transport not connected", transport_type="sse")

        try:
            message_url = urljoin(self.base_url, self.message_path)

            async with self.session.post(
                message_url, json=message, headers={"Content-Type": "application/json"}
            ) as response:
                if response.status not in (200, 201, 202):
                    raise TransportError(
                        f"Message send failed: {response.status}", transport_type="sse"
                    )

            self._update_metrics("messages_sent")
            self._update_metrics("bytes_sent", len(json.dumps(message)))

        except Exception as e:
            self._update_metrics("errors_total")
            raise TransportError(f"Failed to send message: {e}", transport_type="sse")

    async def receive_message(self) -> Dict[str, Any]:
        """Receive message from SSE stream."""
        if not self._connected:
            raise TransportError("Transport not connected", transport_type="sse")

        try:
            # Wait for message from queue
            message = await asyncio.wait_for(
                self._message_queue.get(), timeout=self.timeout
            )

            self._update_metrics("messages_received")
            self._update_metrics("bytes_received", len(json.dumps(message)))

            return message

        except asyncio.TimeoutError:
            raise TransportError("Receive timeout", transport_type="sse")
        except Exception as e:
            self._update_metrics("errors_total")
            raise TransportError(
                f"Failed to receive message: {e}", transport_type="sse"
            )

    async def _read_sse_events(self):
        """Background task to read SSE events."""
        if not self.sse_response:
            return

        try:
            async for line in self.sse_response.content:
                if not self._connected:
                    break

                line_str = line.decode("utf-8").strip()

                # Parse SSE event
                if line_str.startswith("data: "):
                    data_str = line_str[6:]  # Remove "data: " prefix

                    try:
                        message = json.loads(data_str)
                        await self._message_queue.put(message)
                    except json.JSONDecodeError:
                        logger.warning(f"Invalid JSON in SSE event: {data_str}")

        except Exception as e:
            logger.error(f"SSE read error: {e}")
        finally:
            if self._connected:
                await self.disconnect()

    async def _cleanup_connection(self):
        """Clean up connection resources."""
        # Cancel read task
        if self._read_task and not self._read_task.done():
            self._read_task.cancel()
            try:
                await self._read_task
            except asyncio.CancelledError:
                pass

        # Close SSE response
        if self.sse_response:
            self.sse_response.close()
            self.sse_response = None

        # Close session
        if self.session:
            await self.session.close()
            self.session = None


class StreamableHTTPTransport(BaseTransport):
    """StreamableHTTP transport with session management."""

    def __init__(
        self,
        base_url: str,
        session_management: bool = True,
        streaming_threshold: int = 1024,
        chunk_size: int = 8192,
        allow_localhost: bool = False,
        skip_security_validation: bool = False,
        **kwargs,
    ):
        """Initialize StreamableHTTP transport.

        Args:
            base_url: Base URL for the server
            session_management: Enable session management
            streaming_threshold: Size threshold for streaming
            chunk_size: Chunk size for streaming
            allow_localhost: Allow connections to localhost (for testing)
            skip_security_validation: Skip all security validation (for testing)
            **kwargs: Base transport arguments
        """
        super().__init__("streamable_http", **kwargs)

        self.base_url = base_url.rstrip("/")
        self.session_management = session_management
        self.streaming_threshold = streaming_threshold
        self.chunk_size = chunk_size
        self.allow_localhost = allow_localhost
        self.skip_security_validation = skip_security_validation

        # Session state
        self.session: Optional[aiohttp.ClientSession] = None
        self.session_id: Optional[str] = None

    async def connect(self) -> None:
        """Connect and optionally create session."""
        if self._connected:
            return

        # Validate URL (with configurable security)
        if not self.skip_security_validation:
            if not TransportSecurity.validate_url(
                self.base_url, allow_localhost=self.allow_localhost
            ):
                raise TransportError(
                    "Invalid or unsafe URL", transport_type="streamable_http"
                )

        try:
            # Create HTTP session
            timeout = aiohttp.ClientTimeout(total=self.timeout)
            self.session = aiohttp.ClientSession(timeout=timeout)

            # Create server session if enabled
            if self.session_management:
                await self._create_server_session()

            self._connected = True
            self._update_metrics("connections_total")

            logger.info(f"StreamableHTTP transport connected: {self.base_url}")

        except Exception as e:
            self._update_metrics("connections_failed")
            await self._cleanup_connection()
            raise TransportError(
                f"HTTP connection failed: {e}", transport_type="streamable_http"
            )

    async def disconnect(self) -> None:
        """Disconnect and cleanup session."""
        if not self._connected:
            return

        self._connected = False

        # Close server session
        if self.session_management and self.session_id:
            await self._close_server_session()

        await self._cleanup_connection()

        logger.info("StreamableHTTP transport disconnected")

    async def send_message(self, message: Dict[str, Any]) -> None:
        """Send message via HTTP POST."""
        if not self._connected or not self.session:
            raise TransportError(
                "Transport not connected", transport_type="streamable_http"
            )

        try:
            # Prepare URL
            url = urljoin(self.base_url, "/message")

            # Add session ID if using session management
            headers = {"Content-Type": "application/json"}
            if self.session_id:
                headers["X-Session-ID"] = self.session_id

            # Determine if streaming is needed
            message_data = json.dumps(message)
            use_streaming = len(message_data) > self.streaming_threshold

            if use_streaming:
                # Stream large message
                await self._send_streamed_message(url, message_data, headers)
            else:
                # Send normal message
                async with self.session.post(
                    url, json=message, headers=headers
                ) as response:
                    if response.status not in (200, 201, 202):
                        raise TransportError(
                            f"Message send failed: {response.status}",
                            transport_type="streamable_http",
                        )

            self._update_metrics("messages_sent")
            self._update_metrics("bytes_sent", len(message_data))

        except Exception as e:
            self._update_metrics("errors_total")
            raise TransportError(
                f"Failed to send message: {e}", transport_type="streamable_http"
            )

    async def receive_message(self) -> Dict[str, Any]:
        """Receive message via HTTP GET/POST."""
        if not self._connected or not self.session:
            raise TransportError(
                "Transport not connected", transport_type="streamable_http"
            )

        try:
            # Prepare URL
            url = urljoin(self.base_url, "/receive")

            # Add session ID if using session management
            headers = {}
            if self.session_id:
                headers["X-Session-ID"] = self.session_id

            # Receive message
            async with self.session.get(url, headers=headers) as response:
                if response.status == 204:
                    # No message available
                    await asyncio.sleep(0.1)  # Brief delay before retry
                    return await self.receive_message()

                if response.status != 200:
                    raise TransportError(
                        f"Message receive failed: {response.status}",
                        transport_type="streamable_http",
                    )

                # Check if response is streamed
                content_length = response.headers.get("Content-Length")
                if content_length and int(content_length) > self.streaming_threshold:
                    message = await self._receive_streamed_message(response)
                else:
                    message = await response.json()

            self._update_metrics("messages_received")
            self._update_metrics("bytes_received", len(json.dumps(message)))

            return message

        except Exception as e:
            self._update_metrics("errors_total")
            raise TransportError(
                f"Failed to receive message: {e}", transport_type="streamable_http"
            )

    async def _create_server_session(self):
        """Create session with server."""
        if not self.session:
            return

        url = urljoin(self.base_url, "/session")

        async with self.session.post(url) as response:
            if response.status == 201:
                session_data = await response.json()
                self.session_id = session_data.get("session_id")
                logger.info(f"Created server session: {self.session_id}")
            else:
                logger.warning(f"Failed to create server session: {response.status}")

    async def _close_server_session(self):
        """Close session with server."""
        if not self.session or not self.session_id:
            return

        url = urljoin(self.base_url, f"/session/{self.session_id}")

        try:
            async with self.session.delete(url) as response:
                if response.status in (200, 204):
                    logger.info(f"Closed server session: {self.session_id}")
                else:
                    logger.warning(f"Failed to close server session: {response.status}")
        except Exception as e:
            logger.error(f"Error closing server session: {e}")
        finally:
            self.session_id = None

    async def _send_streamed_message(
        self, url: str, message_data: str, headers: Dict[str, str]
    ):
        """Send message using streaming."""
        headers["Transfer-Encoding"] = "chunked"

        async def message_chunks():
            for i in range(0, len(message_data), self.chunk_size):
                yield message_data[i : i + self.chunk_size].encode("utf-8")

        async with self.session.post(
            url, data=message_chunks(), headers=headers
        ) as response:
            if response.status not in (200, 201, 202):
                raise TransportError(
                    f"Streamed message send failed: {response.status}",
                    transport_type="streamable_http",
                )

    async def _receive_streamed_message(
        self, response: aiohttp.ClientResponse
    ) -> Dict[str, Any]:
        """Receive streamed message."""
        chunks = []

        async for chunk in response.content.iter_chunked(self.chunk_size):
            chunks.append(chunk.decode("utf-8"))

        message_data = "".join(chunks)
        return json.loads(message_data)

    async def _cleanup_connection(self):
        """Clean up connection resources."""
        if self.session:
            await self.session.close()
            self.session = None


class WebSocketTransport(BaseTransport):
    """WebSocket transport for real-time communication."""

    def __init__(
        self,
        url: str,
        subprotocols: Optional[List[str]] = None,
        ping_interval: float = 20.0,
        ping_timeout: float = 20.0,
        allow_localhost: bool = False,
        skip_security_validation: bool = False,
        **kwargs,
    ):
        """Initialize WebSocket transport.

        Args:
            url: WebSocket URL
            subprotocols: WebSocket subprotocols
            ping_interval: Ping interval in seconds
            ping_timeout: Ping timeout in seconds
            allow_localhost: Allow connections to localhost (for testing)
            skip_security_validation: Skip all security validation (for testing)
            **kwargs: Base transport arguments
        """
        super().__init__("websocket", **kwargs)

        self.url = url
        self.subprotocols = subprotocols or ["mcp-v1"]
        self.ping_interval = ping_interval
        self.ping_timeout = ping_timeout
        self.allow_localhost = allow_localhost
        self.skip_security_validation = skip_security_validation

        # Connection state
        self.websocket: Optional[websockets.WebSocketServerProtocol] = None
        self._read_task: Optional[asyncio.Task] = None
        self._message_queue: asyncio.Queue = asyncio.Queue()

    async def connect(self) -> None:
        """Connect to WebSocket server."""
        if self._connected:
            return

        # Validate URL (with configurable security)
        if not self.skip_security_validation:
            if not TransportSecurity.validate_url(
                self.url, allow_localhost=self.allow_localhost
            ):
                raise TransportError(
                    "Invalid or unsafe URL", transport_type="websocket"
                )

        try:
            # Connect to WebSocket
            extra_headers = {}
            if self.auth_provider:
                # Add authentication headers
                auth_headers = await self.auth_provider.get_headers()
                extra_headers.update(auth_headers)

            self.websocket = await websockets.connect(
                self.url,
                subprotocols=self.subprotocols,
                extra_headers=extra_headers,
                ping_interval=self.ping_interval,
                ping_timeout=self.ping_timeout,
            )

            # Start reading messages
            self._read_task = asyncio.create_task(self._read_messages())

            self._connected = True
            self._update_metrics("connections_total")

            logger.info(f"WebSocket transport connected: {self.url}")

        except Exception as e:
            self._update_metrics("connections_failed")
            await self._cleanup_connection()
            raise TransportError(
                f"WebSocket connection failed: {e}", transport_type="websocket"
            )

    async def disconnect(self) -> None:
        """Disconnect from WebSocket server."""
        if not self._connected:
            return

        self._connected = False
        await self._cleanup_connection()

        logger.info("WebSocket transport disconnected")

    async def send_message(self, message: Dict[str, Any]) -> None:
        """Send message via WebSocket."""
        if not self._connected or not self.websocket:
            raise TransportError("Transport not connected", transport_type="websocket")

        try:
            message_data = json.dumps(message)
            await self.websocket.send(message_data)

            self._update_metrics("messages_sent")
            self._update_metrics("bytes_sent", len(message_data))

        except Exception as e:
            self._update_metrics("errors_total")
            raise TransportError(
                f"Failed to send message: {e}", transport_type="websocket"
            )

    async def receive_message(self) -> Dict[str, Any]:
        """Receive message from WebSocket."""
        if not self._connected:
            raise TransportError("Transport not connected", transport_type="websocket")

        try:
            message = await asyncio.wait_for(
                self._message_queue.get(), timeout=self.timeout
            )

            self._update_metrics("messages_received")
            self._update_metrics("bytes_received", len(json.dumps(message)))

            return message

        except asyncio.TimeoutError:
            raise TransportError("Receive timeout", transport_type="websocket")
        except Exception as e:
            self._update_metrics("errors_total")
            raise TransportError(
                f"Failed to receive message: {e}", transport_type="websocket"
            )

    async def _read_messages(self):
        """Background task to read WebSocket messages."""
        if not self.websocket:
            return

        try:
            async for message_data in self.websocket:
                if not self._connected:
                    break

                try:
                    message = json.loads(message_data)
                    await self._message_queue.put(message)
                except json.JSONDecodeError:
                    logger.warning(f"Invalid JSON in WebSocket message: {message_data}")

        except websockets.exceptions.ConnectionClosed:
            logger.info("WebSocket connection closed")
        except Exception as e:
            logger.error(f"WebSocket read error: {e}")
        finally:
            if self._connected:
                await self.disconnect()

    async def _cleanup_connection(self):
        """Clean up connection resources."""
        # Cancel read task
        if self._read_task and not self._read_task.done():
            self._read_task.cancel()
            try:
                await self._read_task
            except asyncio.CancelledError:
                pass

        # Close WebSocket
        if self.websocket:
            await self.websocket.close()
            self.websocket = None


class TransportManager:
    """Manager for MCP transport instances."""

    def __init__(self):
        """Initialize transport manager."""
        self._transports: Dict[str, BaseTransport] = {}
        self._transport_factories: Dict[str, Callable] = {
            "stdio": EnhancedStdioTransport,
            "sse": SSETransport,
            "streamable_http": StreamableHTTPTransport,
            "websocket": WebSocketTransport,
        }

    def register_transport_factory(self, transport_type: str, factory: Callable):
        """Register transport factory.

        Args:
            transport_type: Transport type name
            factory: Factory function
        """
        self._transport_factories[transport_type] = factory

    def create_transport(self, transport_type: str, **kwargs) -> BaseTransport:
        """Create transport instance.

        Args:
            transport_type: Transport type
            **kwargs: Transport arguments

        Returns:
            Transport instance
        """
        factory = self._transport_factories.get(transport_type)
        if not factory:
            raise ValueError(f"Unknown transport type: {transport_type}")

        return factory(**kwargs)

    def register_transport(self, name: str, transport: BaseTransport):
        """Register transport instance.

        Args:
            name: Transport name
            transport: Transport instance
        """
        self._transports[name] = transport

    def get_transport(self, name: str) -> Optional[BaseTransport]:
        """Get registered transport.

        Args:
            name: Transport name

        Returns:
            Transport instance or None
        """
        return self._transports.get(name)

    def list_transports(self) -> List[str]:
        """List registered transport names."""
        return list(self._transports.keys())

    async def disconnect_all(self):
        """Disconnect all registered transports."""
        for transport in self._transports.values():
            try:
                await transport.disconnect()
            except Exception as e:
                logger.error(f"Error disconnecting transport: {e}")

        self._transports.clear()


# Global transport manager
_transport_manager: Optional[TransportManager] = None


def get_transport_manager() -> TransportManager:
    """Get global transport manager."""
    global _transport_manager
    if _transport_manager is None:
        _transport_manager = TransportManager()
    return _transport_manager
