"""MCP Transport - 传输层抽象

提供 MCP 协议的传输层实现，支持 Streamable HTTP、SSE 和 Stdio 三种传输方式。
"""

import asyncio
import json
import logging
import os
from abc import ABC, abstractmethod
from typing import Any

import httpx

logger = logging.getLogger(__name__)


class TransportError(Exception):
    """传输层错误"""

    def __init__(self, message: str, cause: Exception | None = None):
        self.cause = cause
        super().__init__(message)


class Transport(ABC):
    """传输层抽象基类

    定义 MCP 协议传输层的统一接口，支持 JSON-RPC 2.0 消息格式。
    """

    @abstractmethod
    async def connect(self) -> None:
        """建立连接"""
        ...

    @abstractmethod
    async def disconnect(self) -> None:
        """关闭连接"""
        ...

    @abstractmethod
    async def send_request(self, method: str, params: dict[str, Any] | None = None) -> Any:
        """发送 JSON-RPC 请求

        Args:
            method: JSON-RPC 方法名
            params: 请求参数

        Returns:
            JSON-RPC 响应的 result 字段
        """
        ...

    @abstractmethod
    async def receive_response(self) -> dict[str, Any]:
        """接收响应

        Returns:
            JSON-RPC 响应消息
        """
        ...


class HTTPTransport(Transport):
    """Streamable HTTP 传输

    使用 httpx.AsyncClient 发送 POST 请求到 MCP 服务器端点，
    支持 JSON-RPC 请求/响应关联。
    """

    def __init__(
        self,
        endpoint: str,
        headers: dict[str, str] | None = None,
        timeout: float = 30.0,
    ):
        self._endpoint = endpoint.rstrip("/")
        self._headers = headers or {}
        self._timeout = timeout
        self._client: httpx.AsyncClient | None = None
        self._request_id = 0
        self._pending: dict[int, asyncio.Future[dict[str, Any]]] = {}

    @property
    def is_connected(self) -> bool:
        return self._client is not None and not self._client.is_closed

    async def connect(self) -> None:
        """建立 HTTP 连接"""
        if self.is_connected:
            return
        self._client = httpx.AsyncClient(
            base_url=self._endpoint,
            headers=self._headers,
            timeout=self._timeout,
        )
        logger.info("HTTPTransport connected to %s", self._endpoint)

    async def disconnect(self) -> None:
        """关闭 HTTP 连接"""
        if self._client is not None and not self._client.is_closed:
            await self._client.aclose()
            self._client = None
            # 取消所有等待中的请求
            for future in self._pending.values():
                if not future.done():
                    future.cancel()
            self._pending.clear()
            logger.info("HTTPTransport disconnected")

    def _next_request_id(self) -> int:
        """生成下一个请求 ID"""
        self._request_id += 1
        return self._request_id

    async def send_request(self, method: str, params: dict[str, Any] | None = None) -> Any:
        """发送 JSON-RPC 请求并等待响应

        Args:
            method: JSON-RPC 方法名
            params: 请求参数

        Returns:
            JSON-RPC 响应的 result 字段

        Raises:
            TransportError: 连接未建立或请求失败
        """
        if not self.is_connected:
            raise TransportError("Transport not connected")

        request_id = self._next_request_id()
        message: dict[str, Any] = {
            "jsonrpc": "2.0",
            "id": request_id,
            "method": method,
        }
        if params is not None:
            message["params"] = params

        try:
            response = await self._client.post(  # type: ignore[union-attr]
                "/",
                json=message,
            )
            response.raise_for_status()
        except httpx.HTTPStatusError as e:
            raise TransportError(f"HTTP error {e.response.status_code}", cause=e) from e
        except httpx.RequestError as e:
            raise TransportError(f"Request failed: {e}", cause=e) from e

        try:
            data = response.json()
        except json.JSONDecodeError as e:
            raise TransportError(f"Invalid JSON response: {e}", cause=e) from e

        # 检查 JSON-RPC 错误
        if "error" in data:
            error = data["error"]
            raise TransportError(
                f"JSON-RPC error {error.get('code')}: {error.get('message')}"
            )

        return data.get("result")

    async def receive_response(self) -> dict[str, Any]:
        """接收响应

        对于 HTTPTransport，响应在 send_request 中同步返回。
        此方法返回最近一次请求的响应（通过内部队列实现）。

        Returns:
            JSON-RPC 响应消息
        """
        if self._pending:
            request_id = next(iter(self._pending))
            future = self._pending.pop(request_id)
            return await future
        raise TransportError("No pending response to receive")


class SSETransport(Transport):
    """Server-Sent Events 传输

    使用 httpx.AsyncClient 连接 SSE 端点接收服务端推送消息，
    通过 HTTP POST 发送客户端请求。
    """

    def __init__(
        self,
        endpoint: str,
        sse_path: str = "/sse",
        message_path: str = "/message",
        headers: dict[str, str] | None = None,
        timeout: float = 30.0,
    ):
        self._endpoint = endpoint.rstrip("/")
        self._sse_path = sse_path
        self._message_path = message_path
        self._headers = headers or {}
        self._timeout = timeout
        self._client: httpx.AsyncClient | None = None
        self._request_id = 0
        self._sse_task: asyncio.Task[None] | None = None
        self._response_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        self._connected = False

    @property
    def is_connected(self) -> bool:
        return self._connected and self._client is not None and not self._client.is_closed

    async def connect(self) -> None:
        """建立 SSE 连接

        连接到 SSE 端点开始监听服务端消息，同时准备 HTTP 客户端用于发送请求。
        """
        if self.is_connected:
            return

        self._client = httpx.AsyncClient(
            base_url=self._endpoint,
            headers=self._headers,
            timeout=self._timeout,
        )
        self._connected = True

        # 启动 SSE 监听任务
        self._sse_task = asyncio.create_task(self._listen_sse())
        logger.info("SSETransport connected to %s", self._endpoint)

    async def disconnect(self) -> None:
        """关闭 SSE 连接"""
        self._connected = False

        if self._sse_task is not None:
            self._sse_task.cancel()
            try:
                await self._sse_task
            except asyncio.CancelledError:
                pass
            self._sse_task = None

        if self._client is not None and not self._client.is_closed:
            await self._client.aclose()
            self._client = None

        # 清空响应队列
        while not self._response_queue.empty():
            self._response_queue.get_nowait()

        logger.info("SSETransport disconnected")

    async def _listen_sse(self) -> None:
        """监听 SSE 事件流"""
        assert self._client is not None
        try:
            async with self._client.stream("GET", self._sse_path) as response:
                response.raise_for_status()
                async for line in response.aiter_lines():
                    if not self._connected:
                        break
                    line = line.strip()
                    if not line or line.startswith(":"):
                        continue
                    if line.startswith("data:"):
                        data_str = line[len("data:"):].strip()
                        try:
                            data = json.loads(data_str)
                            await self._response_queue.put(data)
                        except json.JSONDecodeError:
                            logger.warning("Invalid SSE data: %s", data_str)
        except httpx.RequestError as e:
            if self._connected:
                logger.error("SSE connection error: %s", e)
        except asyncio.CancelledError:
            raise
        except Exception as e:
            if self._connected:
                logger.error("SSE listener error: %s", e)

    def _next_request_id(self) -> int:
        """生成下一个请求 ID"""
        self._request_id += 1
        return self._request_id

    async def send_request(self, method: str, params: dict[str, Any] | None = None) -> Any:
        """通过 HTTP POST 发送 JSON-RPC 请求

        Args:
            method: JSON-RPC 方法名
            params: 请求参数

        Returns:
            JSON-RPC 响应的 result 字段

        Raises:
            TransportError: 连接未建立或请求失败
        """
        if not self.is_connected:
            raise TransportError("Transport not connected")

        request_id = self._next_request_id()
        message: dict[str, Any] = {
            "jsonrpc": "2.0",
            "id": request_id,
            "method": method,
        }
        if params is not None:
            message["params"] = params

        try:
            response = await self._client.post(  # type: ignore[union-attr]
                self._message_path,
                json=message,
            )
            response.raise_for_status()
        except httpx.HTTPStatusError as e:
            raise TransportError(f"HTTP error {e.response.status_code}", cause=e) from e
        except httpx.RequestError as e:
            raise TransportError(f"Request failed: {e}", cause=e) from e

        try:
            data = response.json()
        except json.JSONDecodeError as e:
            raise TransportError(f"Invalid JSON response: {e}", cause=e) from e

        # 检查 JSON-RPC 错误
        if "error" in data:
            error = data["error"]
            raise TransportError(
                f"JSON-RPC error {error.get('code')}: {error.get('message')}"
            )

        return data.get("result")

    async def receive_response(self) -> dict[str, Any]:
        """从 SSE 事件流接收响应

        Returns:
            JSON-RPC 响应消息

        Raises:
            TransportError: 连接未建立
        """
        if not self.is_connected:
            raise TransportError("Transport not connected")

        try:
            return await asyncio.wait_for(
                self._response_queue.get(),
                timeout=self._timeout,
            )
        except asyncio.TimeoutError:
            raise TransportError("Timeout waiting for SSE response")


class StdioTransport(Transport):
    """Stdio 传输

    通过 stdin/stdout 与 MCP Server 子进程通信，使用 newline-delimited JSON-RPC 消息格式。
    """

    def __init__(
        self,
        command: str,
        args: list[str] | None = None,
        env: dict[str, str] | None = None,
        timeout: float = 30.0,
    ):
        self._command = command
        self._args = args or []
        self._env = env
        self._timeout = timeout
        self._process: asyncio.subprocess.Process | None = None
        self._request_id = 0
        self._pending: dict[int, asyncio.Future[Any]] = {}
        self._reader_task: asyncio.Task[None] | None = None
        self._stderr_task: asyncio.Task[None] | None = None
        self._connected = False
        self._notifications: asyncio.Queue[dict[str, Any]] = asyncio.Queue()

    @property
    def is_connected(self) -> bool:
        return (
            self._connected
            and self._process is not None
            and self._process.returncode is None
        )

    def _next_request_id(self) -> int:
        """生成下一个请求 ID"""
        self._request_id += 1
        return self._request_id

    async def connect(self) -> None:
        """启动子进程并完成 MCP 初始化握手

        Raises:
            TransportError: 子进程启动失败或初始化超时
        """
        if self.is_connected:
            return

        # 合并环境变量
        merged_env = dict(os.environ)
        if self._env:
            merged_env.update(self._env)

        try:
            self._process = await asyncio.create_subprocess_exec(
                self._command,
                *self._args,
                env=merged_env,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
        except OSError as e:
            raise TransportError(f"Failed to start process: {self._command}", cause=e) from e

        # 启动 stdout 读取任务
        self._reader_task = asyncio.create_task(self._read_stdout())

        # 启动 stderr 读取任务
        self._stderr_task = asyncio.create_task(self._read_stderr())

        # 发送 initialize 请求并等待响应
        try:
            init_result = await asyncio.wait_for(
                self._send_request_internal(
                    "initialize",
                    {
                        "protocolVersion": "2024-11-05",
                        "capabilities": {},
                        "clientInfo": {"name": "agentkit", "version": "0.1.0"},
                    },
                ),
                timeout=self._timeout,
            )
        except asyncio.TimeoutError:
            await self._cleanup()
            raise TransportError("Timeout waiting for initialize response")
        except TransportError:
            await self._cleanup()
            raise

        # 发送 initialized 通知
        await self._send_notification("notifications/initialized")

        self._connected = True
        logger.info(
            "StdioTransport connected to %s %s",
            self._command,
            " ".join(self._args),
        )

    async def disconnect(self) -> None:
        """关闭子进程连接"""
        self._connected = False
        await self._cleanup()

    async def _cleanup(self) -> None:
        """清理子进程和相关资源"""
        # 取消读取任务
        for task in (self._reader_task, self._stderr_task):
            if task is not None:
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
        self._reader_task = None
        self._stderr_task = None

        # 关闭 stdin
        if self._process is not None and self._process.stdin is not None:
            self._process.stdin.close()
            try:
                await self._process.stdin.drain()
            except Exception:
                pass

        # 等待子进程退出
        if self._process is not None and self._process.returncode is None:
            try:
                await asyncio.wait_for(self._process.wait(), timeout=5.0)
            except asyncio.TimeoutError:
                self._process.kill()
                await self._process.wait()

        self._process = None

        # 取消所有等待中的 Future
        for future in self._pending.values():
            if not future.done():
                future.set_exception(TransportError("Transport disconnected"))
        self._pending.clear()

        logger.info("StdioTransport disconnected")

    async def send_request(self, method: str, params: dict[str, Any] | None = None) -> Any:
        """发送 JSON-RPC 请求并等待响应

        Args:
            method: JSON-RPC 方法名
            params: 请求参数

        Returns:
            JSON-RPC 响应的 result 字段

        Raises:
            TransportError: 连接未建立或请求失败
        """
        if not self.is_connected:
            raise TransportError("Transport not connected")
        return await self._send_request_internal(method, params)

    async def _send_request_internal(
        self, method: str, params: dict[str, Any] | None = None
    ) -> Any:
        """内部请求发送方法（connect 时也可调用）"""
        request_id = self._next_request_id()
        message: dict[str, Any] = {
            "jsonrpc": "2.0",
            "id": request_id,
            "method": method,
        }
        if params is not None:
            message["params"] = params

        await self._write_message(message)

        loop = asyncio.get_running_loop()
        future: asyncio.Future[Any] = loop.create_future()
        self._pending[request_id] = future

        try:
            return await asyncio.wait_for(future, timeout=self._timeout)
        except asyncio.TimeoutError:
            self._pending.pop(request_id, None)
            raise TransportError(f"Timeout waiting for response to {method}")
        except TransportError:
            self._pending.pop(request_id, None)
            raise

    async def _send_notification(self, method: str, params: dict[str, Any] | None = None) -> None:
        """发送 JSON-RPC 通知（无 id，不期待响应）"""
        message: dict[str, Any] = {
            "jsonrpc": "2.0",
            "method": method,
        }
        if params is not None:
            message["params"] = params
        await self._write_message(message)

    async def _write_message(self, message: dict[str, Any]) -> None:
        """将 JSON-RPC 消息写入子进程 stdin"""
        if self._process is None or self._process.stdin is None:
            raise TransportError("Process stdin not available")
        data = (json.dumps(message) + "\n").encode("utf-8")
        self._process.stdin.write(data)
        await self._process.stdin.drain()

    async def receive_response(self) -> dict[str, Any]:
        """接收通知消息

        对于 StdioTransport，请求响应通过 _pending Future 异步返回。
        此方法仅用于获取服务端推送的通知消息。
        空队列时 await 等待（与 SSETransport 行为一致）。

        Returns:
            JSON-RPC 通知消息

        Raises:
            TransportError: 连接未建立或超时
        """
        if not self.is_connected:
            raise TransportError("Transport not connected")

        try:
            return await asyncio.wait_for(
                self._notifications.get(),
                timeout=self._timeout,
            )
        except asyncio.TimeoutError:
            raise TransportError("Timeout waiting for notification")

    async def _read_stdout(self) -> None:
        """持续从子进程 stdout 读取 JSON-RPC 消息"""
        if self._process is None or self._process.stdout is None:
            return

        try:
            while True:
                line = await self._process.stdout.readline()
                if not line:
                    # EOF — 子进程退出
                    if self._connected:
                        logger.warning("StdioTransport: subprocess stdout EOF")
                    break

                line_str = line.decode("utf-8").strip()
                if not line_str:
                    continue

                try:
                    data = json.loads(line_str)
                except json.JSONDecodeError:
                    logger.warning("StdioTransport: invalid JSON from stdout: %s", line_str)
                    continue

                # 响应消息（有 id 字段）
                if "id" in data:
                    request_id = data["id"]
                    future = self._pending.pop(request_id, None)
                    if future is not None and not future.done():
                        if "error" in data:
                            error = data["error"]
                            future.set_exception(
                                TransportError(
                                    f"JSON-RPC error {error.get('code')}: {error.get('message')}"
                                )
                            )
                        else:
                            future.set_result(data.get("result"))
                    elif future is None:
                        logger.warning(
                            "StdioTransport: received response for unknown request id %s",
                            request_id,
                        )

                # 通知消息（有 method 字段，无 id）
                elif "method" in data:
                    await self._notifications.put(data)

        except asyncio.CancelledError:
            raise
        except Exception as e:
            if self._connected:
                logger.error("StdioTransport: stdout reader error: %s", e)

    async def _read_stderr(self) -> None:
        """持续从子进程 stderr 读取并转发到 logger"""
        if self._process is None or self._process.stderr is None:
            return

        try:
            while True:
                line = await self._process.stderr.readline()
                if not line:
                    break
                line_str = line.decode("utf-8", errors="replace").rstrip()
                if line_str:
                    logger.debug("StdioTransport stderr: %s", line_str)
        except asyncio.CancelledError:
            raise
        except Exception as e:
            if self._connected:
                logger.error("StdioTransport: stderr reader error: %s", e)
