"""
"nodelink.transport.websocket
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
WebSocket transport backend. Compatible with websockets v14+.
"""

import asyncio
import json
import logging
import uuid
from typing import Any, Callable, Dict, Optional, Set

import websockets
from websockets.exceptions import ConnectionClosed

logger = logging.getLogger("nodelink.transport.ws")


class WebSocketTransport:
    """Manages all raw WebSocket connections."""

    def __init__(self, host: str, port: int):
        self.host = host
        self.port = port
        self._connections: Dict[str, Any] = {}
        self._on_message: Optional[Callable] = None
        self._on_connect: Optional[Callable] = None
        self._on_disconnect: Optional[Callable] = None
        self._server = None

    def on_message(self, fn: Callable) -> None:
        self._on_message = fn

    def on_connect(self, fn: Callable) -> None:
        self._on_connect = fn

    def on_disconnect(self, fn: Callable) -> None:
        self._on_disconnect = fn

    async def start(self) -> None:
        self._server = await websockets.serve(
            self._handle_connection,
            self.host,
            self.port,
            ping_interval=20,
            ping_timeout=10,
        )
        logger.info(f"Listening on ws://{self.host}:{self.port}")

    async def stop(self) -> None:
        if self._server:
            self._server.close()
            await self._server.wait_closed()

    async def _handle_connection(self, websocket) -> None:
        conn_id = str(uuid.uuid4())
        self._connections[conn_id] = websocket
        logger.debug(f"Connection {conn_id} from {websocket.remote_address}")

        if self._on_connect:
            await self._on_connect(conn_id, websocket.remote_address)

        try:
            async for raw in websocket:
                try:
                    message = json.loads(raw)
                    if self._on_message:
                        await self._on_message(conn_id, message)
                except json.JSONDecodeError:
                    logger.warning(f"Malformed message from {conn_id}: {raw!r}")
                except Exception as e:
                    logger.error(f"Message error on {conn_id}: {e}", exc_info=True)
        except ConnectionClosed as e:
            logger.debug(f"Connection {conn_id} closed: {e.code} {e.reason}")
        except Exception as e:
            logger.error(f"Connection error on {conn_id}: {e}", exc_info=True)
        finally:
            self._connections.pop(conn_id, None)
            if self._on_disconnect:
                await self._on_disconnect(conn_id)

    async def send(self, conn_id: str, event: str, data: Any) -> bool:
        ws = self._connections.get(conn_id)
        if not ws:
            return False
        try:
            await ws.send(json.dumps({"event": event, "data": data}))
            return True
        except ConnectionClosed:
            self._connections.pop(conn_id, None)
            return False
        except Exception as e:
            logger.error(f"Send failed to {conn_id}: {e}")
            return False

    async def broadcast(self, event: str, data: Any, exclude: Optional[Set[str]] = None) -> None:
        exclude = exclude or set()
        payload = json.dumps({"event": event, "data": data})
        snapshot = list(self._connections.items())
        tasks = [(cid, ws.send(payload)) for cid, ws in snapshot if cid not in exclude]
        if not tasks:
            return
        results = await asyncio.gather(*[c for _, c in tasks], return_exceptions=True)
        for (cid, _), result in zip(tasks, results):
            if isinstance(result, Exception):
                self._connections.pop(cid, None)

    @property
    def connection_count(self) -> int:
        return len(self._connections)

    def get_connection_ids(self) -> Set[str]:
        return set(self._connections.keys())
