from __future__ import annotations

import inspect
import ssl
from collections.abc import Callable
from typing import Any

from google.protobuf.message import Message

try:
    from websockets.asyncio.server import serve as ws_serve

    _NEW_WEBSOCKETS_API = True
except ImportError:  # websockets 12.x legacy API
    from websockets import serve as ws_serve

    _NEW_WEBSOCKETS_API = False

from proto_socket.communicator import ParserMap
from proto_socket.ws_client import WsClient


class WsServer:
    def __init__(
        self,
        host: str,
        port: int,
        path: str,
        new_client: Callable[[Any], WsClient] | None = None,
        interval_sec: int = 0,
        wait_sec: int = 0,
        parser_map: ParserMap | None = None,
        ssl_context: ssl.SSLContext | None = None,
    ) -> None:
        self._host = host
        self._port = port
        self._path = path
        self._new_client = new_client
        self._interval_sec = interval_sec
        self._wait_sec = wait_sec
        self._parser_map = parser_map
        self._ssl_context = ssl_context
        self._clients: list[WsClient] = []
        self._server: Any | None = None
        self.on_client_connected: Callable[[WsClient], object] = lambda _: None

    @property
    def port(self) -> int:
        if self._server is None or not self._server.sockets:
            return self._port
        return int(self._server.sockets[0].getsockname()[1])

    def clients(self) -> list[WsClient]:
        return list(self._clients)

    async def start(self) -> None:
        handler = self._handle_websocket if _NEW_WEBSOCKETS_API else self._handle_legacy_websocket
        self._server = await ws_serve(handler, self._host, self._port, ssl=self._ssl_context)

    async def stop(self) -> None:
        if self._server is not None:
            self._server.close()
            await self._server.wait_closed()
            self._server = None
        clients = list(self._clients)
        self._clients.clear()
        for client in clients:
            await client.close()

    async def broadcast(self, message: Message) -> None:
        for client in self.clients():
            await client.communicator.send(message)

    async def _handle_websocket(self, ws: Any) -> None:
        if self._new_client is not None:
            client = self._new_client(ws)
        else:
            if self._parser_map is None:
                raise ValueError("parser_map is required when new_client is not provided")
            client = WsClient(ws, self._interval_sec, self._wait_sec, self._parser_map)

        self._clients.append(client)
        client.add_disconnect_listener(lambda c: self._remove_client(c))
        result = self.on_client_connected(client)
        if inspect.isawaitable(result):
            await result
        await client.wait_closed()

    def _remove_client(self, client: WsClient) -> None:
        if client in self._clients:
            self._clients.remove(client)

    async def _handle_legacy_websocket(self, ws: Any, _path: str) -> None:
        await self._handle_websocket(ws)