from __future__ import annotations
import logging, asyncio, json
from typing import TYPE_CHECKING, Any, Callable, Awaitable, Self

if TYPE_CHECKING:
    import websockets

from ._core import JsonRpcConnection
from ._exceptions import *

logger = logging.getLogger(__name__)

class JsonRpcWebSocketConnection(JsonRpcConnection):
    def __init__(self, ws: websockets.ClientConnection):
        super().__init__()
        self._ws = ws

    @classmethod
    async def connect(cls, uri: str, **kwargs) -> Self:
        """
        Client-side: connects to a JSON-RPC server over WebSocket.

        Args:
            uri: The WebSocket URI to connect to.
        """
        import websockets
        headers = kwargs.pop("headers", {})
        logger.info("Connecting to JSON-RPC WebSocket server at %s", uri)
        ws = await websockets.connect(
            uri,
            additional_headers=headers,
        )
        connection = cls(ws, **kwargs)
        connection._start_read_loop()
        return connection

    @classmethod
    async def serve(cls, host: str, port: int, *, handler: Callable[[Self], Awaitable] | None = None, **kwargs):
        """
        Server-side: starts a JSON-RPC WebSocket server.

        Args:
            host: The host to bind the server to.
            port: The port to bind the server to.
            handler: Async function called with each new connection.
                Should await the 'wait_closed' method to keep the connection alive for the
                duration of the session.
        """
        if handler is None:
            async def default_handler(conn: Self):
                await conn.wait_closed()
            handler = default_handler
        import websockets
        logger.info("Starting JSON-RPC WebSocket server on %s:%s", host, port)
        async def on_connect(ws):
            logger.debug("New connection from %s", ws.remote_address)
            conn = cls(ws, **kwargs)
            conn._start_read_loop()
            await handler(conn)

        async with websockets.serve(on_connect, host, port):
            await asyncio.Future()

    async def _send_bytes(self, data: bytes) -> None:
        await self._ws.send(data)

    async def _cleanup(self) -> None:
        import websockets
        try:
            await self._ws.close()
        except websockets.ConnectionClosed:
            pass

    async def _read_loop(self) -> None:
        import websockets
        try:
            async for msg in self._ws:
                await self._handle_incoming_message(json.loads(msg))
        except websockets.ConnectionClosed as e:
            logger.debug("WebSocket connection closed")
        except Exception:
            logger.error("Error in WebSocket read loop", exc_info=True)
        finally:
            await self.close()


__all__ = ["JsonRpcWebSocketConnection"]
