"""WebSocket transport for AICC: client and server."""

from __future__ import annotations

import asyncio
import json
from typing import Any

from aicc.bridge import Bridge
from aicc.errors import ConnectionError_


class WebSocketClientTransport:
    """AICC client transport over WebSocket.

    Example:
        t = WebSocketClientTransport("ws://localhost:8765")
        client = AICCClient(t)
    """

    def __init__(self, uri: str):
        self._uri = uri
        self._ws = None  # type: ignore[var-annotated]

    async def connect(self) -> None:
        try:
            import websockets  # type: ignore
        except ImportError as e:  # pragma: no cover
            raise ConnectionError_(
                "websockets is required for WebSocketClientTransport. "
                "Install with: pip install aicc[ws] or pip install websockets"
            ) from e
        self._ws = await websockets.connect(self._uri, max_size=64 * 1024 * 1024)

    async def close(self) -> None:
        if self._ws is not None:
            await self._ws.close()
            self._ws = None

    async def send(self, message: dict[str, Any]) -> None:
        if self._ws is None:
            raise ConnectionError_("WebSocket not connected.")
        await self._ws.send(json.dumps(message))

    async def receive(self) -> dict[str, Any]:
        if self._ws is None:
            raise ConnectionError_("WebSocket not connected.")
        raw = await self._ws.recv()
        if isinstance(raw, bytes):
            raw = raw.decode("utf-8")
        return json.loads(raw)


class WebSocketServer:
    """AICC server transport: serves a Bridge over WebSocket.

    Example:
        bridge = Bridge(name="room")
        async with WebSocketServer(bridge, port=8765):
            await asyncio.Future()  # run forever
    """

    def __init__(self, bridge: Bridge, host: str = "127.0.0.1", port: int = 8765):
        self._bridge = bridge
        self._host = host
        self._port = port
        self._server = None  # type: ignore[var-annotated]
        self._connections: set[Any] = set()

    async def __aenter__(self) -> "WebSocketServer":
        await self.start()
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        await self.stop()

    async def start(self) -> None:
        try:
            import websockets  # type: ignore
        except ImportError as e:  # pragma: no cover
            raise ConnectionError_(
                "websockets is required for WebSocketServer. "
                "Install with: pip install websockets"
            ) from e

        async def handler(ws) -> None:  # type: ignore[no-untyped-def]
            self._connections.add(ws)
            session = self._bridge.open_session()
            try:
                manifest = session.manifest
                if manifest is not None:
                    await ws.send(json.dumps(manifest.model_dump(by_alias=True)))
                async for raw in ws:
                    if isinstance(raw, bytes):
                        raw = raw.decode("utf-8")
                    msg = json.loads(raw)
                    response = await self._bridge.handle_message(session.session_id, msg)
                    if response is not None:
                        await ws.send(json.dumps(response))
            finally:
                self._bridge.close_session(session.session_id)
                self._connections.discard(ws)

        self._server = await websockets.serve(handler, self._host, self._port)

    async def stop(self) -> None:
        if self._server is not None:
            self._server.close()
            await self._server.wait_closed()
            self._server = None
        for ws in list(self._connections):
            try:
                await ws.close()
            except Exception:  # noqa: BLE001
                pass
        self._connections.clear()