osint-dashboard/app/ws_manager.py

"""In-memory WebSocket pub/sub with viewport filtering.

Zero extra deps. Ingest workers publish AIS/ADS-B points; only clients whose
current map bbox contains the point receive the payload. No Redis/Kafka.
"""

from __future__ import annotations

import asyncio
from typing import Any

BBox = tuple[float, float, float, float]  # minlon, minlat, maxlon, maxlat


def point_in_bbox(lon: float, lat: float, bbox: BBox | None) -> bool:
    """True if (lon, lat) sits inside an axis-aligned viewport."""
    if bbox is None:
        return False
    minlon, minlat, maxlon, maxlat = bbox
    return minlon <= lon <= maxlon and minlat <= lat <= maxlat


class ConnectionManager:
    """Maps Tailscale/browser clients → viewport bbox + per-client queue."""

    def __init__(self) -> None:
        self._queues: dict[str, asyncio.Queue] = {}
        self._viewports: dict[str, BBox] = {}

    def register(self, client_id: str, maxsize: int = 256) -> asyncio.Queue:
        q: asyncio.Queue = asyncio.Queue(maxsize=maxsize)
        self._queues[client_id] = q
        return q

    def unregister(self, client_id: str) -> None:
        self._queues.pop(client_id, None)
        self._viewports.pop(client_id, None)

    def set_viewport(self, client_id: str, bbox: BBox) -> None:
        if client_id in self._queues:
            self._viewports[client_id] = bbox

    def viewport_of(self, client_id: str) -> BBox | None:
        return self._viewports.get(client_id)

    def viewports(self) -> list[BBox]:
        return list(self._viewports.values())

    def has_clients(self) -> bool:
        return bool(self._queues)

    async def publish_point(
        self,
        kind: str,
        payload: dict[str, Any],
        *,
        lat: float,
        lon: float,
    ) -> int:
        """Enqueue `{type, payload}` for clients whose viewport contains the point.

        Drops the oldest queued message if a client's buffer is full so a slow
        tab cannot stall ingest. Returns the number of clients that got a copy.
        """
        msg = {"type": kind, "payload": payload}
        sent = 0
        for client_id, queue in list(self._queues.items()):
            if not point_in_bbox(lon, lat, self._viewports.get(client_id)):
                continue
            if queue.full():
                try:
                    queue.get_nowait()
                except asyncio.QueueEmpty:
                    pass
            try:
                queue.put_nowait(msg)
            except asyncio.QueueFull:
                continue
            sent += 1
        return sent


manager = ConnectionManager()