"""Transport layer implementations for MCP."""

import asyncio
from enum import Enum
from typing import Any

import httpx
from mcp.client.stdio import stdio_client

from mcp import StdioServerParameters

from ..core.logger import get_logger

logger = get_logger(__name__)


class TransportType(Enum):
    """Supported transport types."""

    STDIO = "stdio"
    SSE = "sse"
    HTTP = "http"


class SSEClientTransport:
    """SSE client transport for MCP."""

    def __init__(self, url: str, headers: dict[str, str] | None = None):
        """Initialize SSE transport.

        Args:
            url: SSE server URL
            headers: Optional HTTP headers
        """
        self.url = url.rstrip("/")
        self.headers = headers or {}
        self._client: httpx.AsyncClient | None = None
        self._read_stream: Any = None
        self._write_stream: Any = None

    async def __aenter__(self):
        """Enter context manager.

        Returns:
            Tuple of (read_stream, write_stream)
        """
        try:
            from mcp.client.sse import sse_client
        except ImportError:
            raise RuntimeError(
                "SSE transport requires 'mcp[sse]'. Install with: pip install 'mcp[sse]'"
            )

        logger.info(f"Connecting to SSE server at {self.url}")

        # 使用 MCP SDK 的 sse_client
        self._client = sse_client(self.url, self.headers)

        streams = await self._client.__aenter__()
        self._read_stream, self._write_stream = streams

        logger.info("SSE connection established")

        return streams

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Exit context manager.

        Args:
            exc_type: Exception type
            exc_val: Exception value
            exc_tb: Exception traceback
        """
        if self._client:
            await self._client.__aexit__(exc_type, exc_val, exc_tb)
            logger.info("SSE connection closed")


class HTTPClientTransport:
    """HTTP client transport for MCP (JSON-RPC over HTTP)."""

    def __init__(
        self,
        url: str,
        headers: dict[str, str] | None = None,
        timeout: float = 30.0,
    ):
        """Initialize HTTP transport.

        Args:
            url: HTTP server URL
            headers: Optional HTTP headers
            timeout: Request timeout in seconds
        """
        self.url = url.rstrip("/")
        self.headers = headers or {}
        self.timeout = timeout
        self._client: httpx.AsyncClient | None = None
        self._session_initialized = False
        self._request_id = 0
        self._lock = asyncio.Lock()

        # 创建读写流对象
        self._read_stream = self._create_read_stream()
        self._write_stream = self._create_write_stream()

    def _create_read_stream(self):
        """Create read stream for HTTP transport.

        Returns:
            Read stream object with recv() method
        """
        queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()

        class HTTPReadStream:
            """HTTP read stream."""

            def __init__(self, queue_ref: asyncio.Queue, parent: "HTTPClientTransport"):
                self.queue = queue_ref
                self.parent = parent

            async def recv(self) -> dict[str, Any]:
                """Receive a message from the stream.

                Returns:
                    Message dictionary
                """
                return await self.queue.get()

        return HTTPReadStream(queue, self)

    def _create_write_stream(self):
        """Create write stream for HTTP transport.

        Returns:
            Write stream object with send() method
        """

        class HTTPWriteStream:
            """HTTP write stream."""

            def __init__(self, parent: "HTTPClientTransport"):
                self.parent = parent

            async def send(self, message: dict[str, Any]) -> None:
                """Send a message through the stream.

                Args:
                    message: Message dictionary to send
                """
                await self.parent._send_http_request(message)

        return HTTPWriteStream(self)

    async def _send_http_request(self, message: dict[str, Any]) -> dict[str, Any]:
        """Send HTTP request and get response.

        Args:
            message: JSON-RPC message

        Returns:
            Response dictionary

        Raises:
            RuntimeError: If client is not initialized
        """
        if not self._client:
            raise RuntimeError("HTTP client is not initialized")

        method = message.get("method")
        params = message.get("params", {})
        req_id = message.get("id")

        logger.debug(f"HTTP MCP request: method={method}, id={req_id}")

        # 处理 initialize 方法
        if method == "initialize":
            self._session_initialized = True

        # 发送 HTTP POST 请求
        response = await self._client.post(
            f"{self.url}/mcp",
            json={"jsonrpc": "2.0", "id": req_id, "method": method, "params": params},
            headers=self.headers,
            timeout=self.timeout,
        )

        response.raise_for_status()
        result = response.json()

        logger.debug(f"HTTP MCP response: id={req_id}")

        return result

    async def __aenter__(self):
        """Enter context manager.

        Returns:
            Tuple of (read_stream, write_stream)
        """
        logger.info(f"Connecting to HTTP server at {self.url}")

        # 创建 HTTP 客户端
        self._client = httpx.AsyncClient(timeout=self.timeout, headers=self.headers)

        # 初始化会话
        await self._initialize_session()

        logger.info("HTTP connection established")

        return (self._read_stream, self._write_stream)

    async def _initialize_session(self):
        """Initialize MCP session via HTTP.

        Raises:
            RuntimeError: If initialization fails
        """
        try:
            response = await self._client.post(
                f"{self.url}/mcp",
                json={
                    "jsonrpc": "2.0",
                    "id": 0,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2024-11-05",
                        "capabilities": {
                            "tools": {},
                            "resources": {},
                            "prompts": {},
                        },
                        "clientInfo": {"name": "mcp-router", "version": "1.0.0"},
                    },
                },
            )

            response.raise_for_status()
            result = response.json()

            if "error" in result:
                raise RuntimeError(f"HTTP initialization failed: {result['error']}")

            if "result" not in result:
                raise RuntimeError("HTTP initialization failed: no result in response")

            logger.info("HTTP MCP session initialized successfully")

        except Exception as e:
            logger.error(f"Failed to initialize HTTP session: {e}")
            raise

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Exit context manager.

        Args:
            exc_type: Exception type
            exc_val: Exception value
            exc_tb: Exception traceback
        """
        if self._client:
            await self._client.aclose()
            logger.info("HTTP connection closed")


def create_transport(
    transport_type: str,
    command: str,
    args: list[str] | None = None,
    env: dict[str, str] | None = None,
    url: str | None = None,
    headers: dict[str, str] | None = None,
) -> Any:
    """Create appropriate transport based on type.

    Args:
        transport_type: Type of transport (stdio, sse, http)
        command: Command to execute (for stdio)
        args: Command arguments (for stdio)
        env: Environment variables (for stdio)
        url: Server URL (for sse/http)
        headers: HTTP headers (for sse/http)

    Returns:
        Transport context manager

    Raises:
        ValueError: If transport type is not supported or required params missing
    """
    args = args or []
    env = env or {}

    if transport_type == TransportType.STDIO.value:
        server_params = StdioServerParameters(command=command, args=args, env=env)
        return stdio_client(server_params)

    elif transport_type == TransportType.SSE.value:
        if not url:
            raise ValueError("URL is required for SSE transport")
        return SSEClientTransport(url, headers)

    elif transport_type == TransportType.HTTP.value:
        if not url:
            raise ValueError("URL is required for HTTP transport")
        return HTTPClientTransport(url, headers)

    else:
        raise ValueError(f"Unsupported transport type: {transport_type}")
