"""Local WebSocket server for AiCippy CLI to Browser Extension communication."""

from __future__ import annotations

import asyncio
import json
import logging
import threading
from collections.abc import Callable
from typing import Any

import websockets
from websockets.asyncio.server import ServerConnection

logger = logging.getLogger(__name__)


class LocalServerBindError(RuntimeError):
    """Exception raised when local server fails to bind."""

    def __init__(self, port: int) -> None:
        """Initialize the exception."""
        super().__init__(f"Could not bind to any port in range {port}-{port+10}")

class BrowserServer:
    """Local WebSocket server for communicating with the browser extension.
    Runs on localhost:9000.
    """  # noqa: D205

    def __init__(self, host: str = "localhost", port: int = 9000) -> None:  # noqa: D107
        self.host = host
        self.port = port
        self.clients: set[ServerConnection] = set()
        self.loop: asyncio.AbstractEventLoop | None = None
        self.server: Any = None
        self._stop_event = asyncio.Event()
        self._handlers: dict[str, list[Callable[[dict[str, Any]], Any]]] = {}

    async def _handler(self, websocket: ServerConnection) -> None:
        """Handle new WebSocket connections."""
        self.clients.add(websocket)
        logger.info(f"Browser extension connected: {websocket.remote_address}")
        try:
            async for message in websocket:
                try:
                    data = json.loads(message)
                    msg_type = data.get("type")
                    if msg_type in self._handlers:
                        for handler in self._handlers[msg_type]:
                            if asyncio.iscoroutinefunction(handler):
                                await handler(data)
                            else:
                                handler(data)
                except json.JSONDecodeError:
                    logger.warning(f"Invalid JSON from browser: {message!r}")
        except websockets.exceptions.ConnectionClosed:
            pass
        finally:
            self.clients.remove(websocket)
            logger.info(f"Browser extension disconnected: {websocket.remote_address}")

    def on(self, msg_type: str, handler: Callable[[dict[str, Any]], Any]) -> None:
        """Register a handler for a message type."""
        if msg_type not in self._handlers:
            self._handlers[msg_type] = []
        self._handlers[msg_type].append(handler)

    async def broadcast(self, message: dict[str, Any]) -> None:
        """Broadcast a message to all connected extension clients."""
        if not self.clients:
            return
        msg_str = json.dumps(message)
        await asyncio.gather(
            *[client.send(msg_str) for client in self.clients],
            return_exceptions=True,
        )

    async def start(self) -> None:
        """Start the WebSocket server."""
        self.loop = asyncio.get_running_loop()
        
        for port in range(self.port, self.port + 11):
            try:
                self.server = await websockets.asyncio.server.serve(self._handler, self.host, port)  # type: ignore[attr-defined]  # websockets stubs incomplete
                self.port = port
                break
            except OSError:
                continue
        else:
            raise LocalServerBindError(self.port)
            
        logger.info(f"Local browser server started on ws://{self.host}:{self.port}")
        await self._stop_event.wait()
        self.server.close()
        await self.server.wait_closed()

    def stop(self) -> None:
        """Stop the server."""
        if self.loop:
            self.loop.call_soon_threadsafe(self._stop_event.set)


def start_server_in_thread(server: BrowserServer) -> threading.Thread:
    """Start the server in a separate thread."""

    def run() -> None:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        loop.run_until_complete(server.start())

    thread = threading.Thread(target=run, daemon=True)
    thread.start()
    return thread
