"""Local WebSocket server for direct browser connections on the same machine."""

import asyncio
import json
import logging
import secrets
from typing import Optional, Set, Dict
from urllib.parse import urlparse, parse_qs

from websockets.asyncio.server import serve, ServerConnection

from .message_handler import MessageHandler
from .rate_limiter import SlidingWindowRateLimiter

logger = logging.getLogger("skop-runner")

# Rate limiting for incoming messages (prevents DoS from local processes)
MAX_INCOMING_PER_SECOND = 200
INCOMING_WINDOW_SECONDS = 1.0


class LocalWebSocketServer:
    """
    Local WebSocket server that accepts direct browser connections.

    Binds to 127.0.0.1 on an OS-assigned port. Connections must provide a
    token via ?token=... query parameter to prevent Cross-Site WebSocket
    Hijacking (CSWSH). The token is generated at startup and advertised
    to the frontend via the authenticated runner API.

    The _broadcast_notification method is registered with the
    NotificationMultiplexer so all connected local clients receive
    kernel output, status changes, and other notifications.
    No outgoing rate limiting — per-client send locks and TCP
    backpressure naturally throttle the caller.
    """

    def __init__(self, message_handler: MessageHandler):
        self._message_handler = message_handler
        self._server = None
        self._port: Optional[int] = None
        self._token: str = secrets.token_urlsafe(32)
        self._clients: Set[ServerConnection] = set()
        self._client_locks: Dict[int, asyncio.Lock] = {}

    @property
    def port(self) -> Optional[int]:
        """The OS-assigned port, or None if not yet started."""
        return self._port

    @property
    def token(self) -> str:
        """The connection token. Advertised via heartbeat metadata."""
        return self._token

    async def start(self) -> int:
        """
        Bind the server and start accepting connections.

        Returns the OS-assigned port number. Must be called before the first
        heartbeat so the port can be advertised in runner metadata.
        """
        self._server = await serve(
            self._handle_connection,
            "127.0.0.1",
            0,
            ping_interval=20,
            ping_timeout=10,
            max_size=50 * 1024 * 1024,  # 50MB — match file write limit
        )

        sockets = self._server.sockets
        self._port = sockets[0].getsockname()[1]
        logger.info(f"Local WebSocket server listening on ws://127.0.0.1:{self._port}")
        return self._port

    async def stop(self) -> None:
        """Shut down the server and close all client connections."""
        if self._server:
            self._server.close()
            await self._server.wait_closed()
            self._server = None
            self._port = None
        self._clients.clear()
        self._client_locks.clear()
        logger.info("Local WebSocket server stopped")

    async def _handle_connection(self, ws: ServerConnection) -> None:
        """Handle a single browser client connection lifecycle."""
        # Validate connection token (prevents CSWSH attacks)
        path = ws.request.path if ws.request else ""
        params = parse_qs(urlparse(path).query)
        provided = params.get("token", [None])[0]
        if not provided or not secrets.compare_digest(provided, self._token):
            await ws.close(4003, "Invalid token")
            return

        client_id = id(ws)
        self._clients.add(ws)
        self._client_locks[client_id] = asyncio.Lock()
        peer = ws.remote_address
        incoming_limiter = SlidingWindowRateLimiter(
            MAX_INCOMING_PER_SECOND, INCOMING_WINDOW_SECONDS,
            label="incoming local WS messages",
        )
        pending_tasks: set[asyncio.Task] = set()
        logger.info(
            f"Local WS client connected from {peer} (total: {len(self._clients)})"
        )

        try:
            async for raw_message in ws:
                # Rate limit incoming messages to prevent DoS
                if not incoming_limiter.allow():
                    continue

                try:
                    message = json.loads(raw_message)
                except json.JSONDecodeError:
                    logger.warning("Local WS: invalid JSON received, ignoring")
                    continue

                # Handle concurrently so long-running RPCs don't block others
                task = asyncio.create_task(
                    self._handle_and_respond(ws, message, client_id)
                )
                pending_tasks.add(task)
                task.add_done_callback(pending_tasks.discard)
        except Exception as e:
            logger.debug(f"Local WS client error: {e}")
        finally:
            self._clients.discard(ws)
            self._client_locks.pop(client_id, None)
            # Drain in-flight handlers, then cancel any that take too long
            if pending_tasks:
                done, remaining = await asyncio.wait(
                    pending_tasks, timeout=2.0
                )
                for task in remaining:
                    task.cancel()
                if remaining:
                    await asyncio.gather(*remaining, return_exceptions=True)
            logger.info(
                f"Local WS client disconnected (total: {len(self._clients)})"
            )

    async def _handle_and_respond(
        self,
        ws: ServerConnection,
        message: dict,
        client_id: int,
    ) -> None:
        """Route message through MessageHandler and send response to the caller."""
        try:
            response = await self._message_handler.handle_message(message)
            lock = self._client_locks.get(client_id)
            if lock and ws in self._clients:
                async with lock:
                    await ws.send(json.dumps(response))
        except Exception as e:
            logger.error(f"Local WS: error handling message: {e}", exc_info=True)

    async def _broadcast_notification(self, method: str, params: dict) -> None:
        """
        Send a JSON-RPC notification to all connected local WS clients.

        This is the NotificationSender registered with the multiplexer.
        Sends to all clients concurrently so a slow client doesn't block others.
        """
        if not self._clients:
            return

        notification = json.dumps({
            "jsonrpc": "2.0",
            "method": method,
            "params": params,
        })

        # Snapshot to avoid mutation during iteration
        clients = set(self._clients)
        await asyncio.gather(
            *[self._send_to_client(c, id(c), notification) for c in clients],
            return_exceptions=True,
        )

    async def _send_to_client(
        self, client: ServerConnection, client_id: int, data: str
    ) -> None:
        """Send a message to a single client, removing it on failure."""
        lock = self._client_locks.get(client_id)
        if lock is None:
            return
        try:
            async with lock:
                await client.send(data)
        except Exception:
            logger.debug(f"Local WS: failed to send to client {client_id}")
            self._clients.discard(client)
            self._client_locks.pop(client_id, None)
