"""WebSocket transport for MCP servers.

This transport uses the `websockets` dependency, but we import it lazily so
`mcpbridge` can be imported in environments where WebSockets support isn't
installed (e.g. during unit tests that don't exercise the transport).
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any

try:
    import websockets  # type: ignore
    from websockets.legacy.client import WebSocketClientProtocol  # type: ignore
    from websockets.exceptions import ConnectionClosed  # type: ignore
except ImportError:  # pragma: no cover
    websockets = None  # type: ignore[assignment]
    WebSocketClientProtocol = Any  # type: ignore[misc,assignment]
    ConnectionClosed = Exception  # type: ignore[misc,assignment]

from mcpbridge.exceptions import MCPProtocolError, TransportConnectionError, TransportDisconnectedError
from mcpbridge.transport.base import BaseTransport

logger = logging.getLogger("mcpbridge")


class WebSocketTransport(BaseTransport):
    """WebSocket MCP transport implementation."""

    def __init__(
        self,
        url: str,
        headers: dict[str, str] | None = None,
        timeout: int = 30,
        ping_interval: int = 30,
        ping_timeout: int = 20,
    ) -> None:
        self._url = url
        self._headers = headers or {}
        self._timeout = timeout
        self._ping_interval = ping_interval
        self._ping_timeout = ping_timeout

        self._ws: WebSocketClientProtocol | None = None
        self._recv_task: asyncio.Task[None] | None = None
        self._connected = False
        self._write_lock = asyncio.Lock()
        self._pending: dict[str | int, asyncio.Future[dict[str, Any]]] = {}

        self.notifications_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()

    @property
    def is_connected(self) -> bool:
        return self._connected

    async def connect(self) -> None:
        if self._connected:
            return
        if websockets is None:
            raise TransportConnectionError(
                "Missing `websockets` dependency required for WebSocket MCP transport. "
                "Install it (or install mcpbridge) to use this transport."
            )

        backoff = 0.5
        for _ in range(6):
            try:
                self._ws = await websockets.connect(
                    self._url,
                    extra_headers=self._headers,
                    ping_interval=self._ping_interval,
                    ping_timeout=self._ping_timeout,
                )
                self._connected = True
                self._recv_task = asyncio.create_task(self._recv_loop(), name="mcpbridge-ws-recv")
                return
            except Exception as e:
                logger.warning("WebSocket connect failed (%s). Retrying in %.1fs...", e, backoff)
                await asyncio.sleep(backoff)
                backoff = min(backoff * 2, 8.0)

        raise TransportConnectionError("Unable to connect to WebSocket MCP server")

    async def close(self) -> None:
        self._connected = False
        if self._recv_task:
            self._recv_task.cancel()
        if self._ws:
            try:
                await self._ws.close()
            except Exception:
                pass
        self._ws = None
        for fut in list(self._pending.values()):
            if not fut.done():
                fut.set_exception(TransportDisconnectedError("Transport disconnected"))
        self._pending.clear()

    async def send_notification(self, method: str, params: dict[str, Any]) -> None:
        msg: dict[str, Any] = {"jsonrpc": "2.0", "method": method}
        if params is not None:
            msg["params"] = params
        await self._send_message(msg)

    async def send(self, message: dict[str, Any]) -> dict[str, Any]:
        if not self._connected:
            await self.connect()
        assert self._ws is not None

        msg_id = message.get("id")
        if msg_id is None:
            raise MCPProtocolError("send() requires a JSON-RPC `id` field")

        loop = asyncio.get_running_loop()
        fut: asyncio.Future[dict[str, Any]] = loop.create_future()
        self._pending[msg_id] = fut
        try:
            await self._send_message(message)
            resp = await asyncio.wait_for(fut, timeout=self._timeout)
            return resp
        except asyncio.TimeoutError as e:
            self._pending.pop(msg_id, None)
            raise TimeoutError("Transport timed out") from e
        finally:
            self._pending.pop(msg_id, None)

    async def _send_message(self, msg: dict[str, Any]) -> None:
        if not self._connected or not self._ws:
            raise TransportDisconnectedError("Transport not connected")
        async with self._write_lock:
            await self._ws.send(json.dumps(msg, ensure_ascii=False))

    async def _recv_loop(self) -> None:
        assert self._ws is not None
        try:
            async for raw in self._ws:
                try:
                    msg = json.loads(raw)
                except Exception as e:
                    raise MCPProtocolError("Malformed JSON from MCP WebSocket") from e

                msg_id = msg.get("id")
                method = msg.get("method")
                if msg_id is not None and msg_id in self._pending and not self._pending[msg_id].done():
                    self._pending[msg_id].set_result(msg)
                elif msg_id is None and isinstance(method, str):
                    await self.notifications_queue.put(msg)
        except ConnectionClosed as e:
            self._connected = False
            for fut in list(self._pending.values()):
                if not fut.done():
                    fut.set_exception(TransportDisconnectedError(f"WebSocket disconnected: {e}"))
            self._pending.clear()
            logger.warning("WebSocket receive loop closed: %s", e)
        except asyncio.CancelledError:
            return
        except Exception as e:
            self._connected = False
            for fut in list(self._pending.values()):
                if not fut.done():
                    fut.set_exception(TransportDisconnectedError(f"WebSocket receive error: {e}"))
            self._pending.clear()
            logger.error("WebSocket receive loop error: %s", e)

