planet/backend/app/core/websocket/broadcaster.py

"""Data broadcaster for WebSocket connections"""

import asyncio
from datetime import UTC, datetime
from typing import Dict, Any

from app.core.logging import get_logger
from app.core.time import to_iso8601_utc
from app.core.websocket.manager import manager


EARTH_UPDATES_CHANNEL = "earth_updates"
VESSEL_STATE_QUERY_BATCH_SIZE = 1000
logger = get_logger(__name__, service="websocket")


class DataBroadcaster:
    """Periodically broadcasts data to connected WebSocket clients"""

    def __init__(self):
        self.running = False
        self.tasks: Dict[str, asyncio.Task] = {}
        self._pending_vessel_updates: Dict[str, Dict[str, Any]] = {}
        self._vessel_flush_interval = 1.0

    async def get_dashboard_stats(self) -> Dict[str, Any]:
        """Get dashboard statistics"""
        return {
            "total_datasources": 9,
            "active_datasources": 8,
            "tasks_today": 45,
            "success_rate": 97.8,
            "last_updated": to_iso8601_utc(datetime.now(UTC)),
            "alerts": {"critical": 0, "warning": 2, "info": 5},
        }

    async def broadcast_stats(self, interval: int = 5):
        """Broadcast dashboard stats periodically"""
        while self.running:
            try:
                stats = await self.get_dashboard_stats()
                await manager.broadcast(
                    {
                        "type": "data_frame",
                        "channel": "dashboard",
                        "timestamp": to_iso8601_utc(datetime.now(UTC)),
                        "payload": {"stats": stats},
                    },
                    channel="dashboard",
                )
            except Exception:
                pass
            await asyncio.sleep(interval)

    async def broadcast_alert(self, alert: Dict[str, Any]):
        """Broadcast an alert to all connected clients"""
        await manager.broadcast(
            {
                "type": "alert_notification",
                "timestamp": to_iso8601_utc(datetime.now(UTC)),
                "data": {"alert": alert},
            }
        )

    async def broadcast_gpu_update(self, data: Dict[str, Any]):
        """Broadcast GPU cluster update"""
        await manager.broadcast(
            {
                "type": "data_frame",
                "channel": "gpu_clusters",
                "timestamp": to_iso8601_utc(datetime.now(UTC)),
                "payload": data,
            }
        )

    async def broadcast_custom(self, channel: str, data: Dict[str, Any]):
        """Broadcast custom data to a specific channel"""
        if channel == "vessels":
            self.enqueue_vessel_update(data)
            return
        await manager.broadcast(
            {
                "type": "data_frame",
                "channel": channel,
                "timestamp": to_iso8601_utc(datetime.now(UTC)),
                "payload": data,
            },
            channel=channel,
        )

    async def broadcast_earth_update(self, data: Dict[str, Any]):
        """Broadcast Earth visualization refresh hints to connected clients."""
        await self.broadcast_custom(EARTH_UPDATES_CHANNEL, data)

    def enqueue_vessel_update(self, data: Dict[str, Any]):
        vessels = data.get("vessels") if isinstance(data, dict) else None
        if not isinstance(vessels, list):
            return
        source = data.get("source")
        action = data.get("action") or "upsert"
        created = data.get("created")
        for vessel in vessels:
            if not isinstance(vessel, dict):
                continue
            mmsi = vessel.get("mmsi")
            if mmsi in (None, ""):
                continue
            self._pending_vessel_updates[str(mmsi)] = {
                **vessel,
                "_source": source,
                "_action": action,
                "_created": created,
            }

    async def flush_vessel_updates(self):
        if not self._pending_vessel_updates:
            return
        pending = self._pending_vessel_updates
        self._pending_vessel_updates = {}
        try:
            vessels = await self._load_current_vessel_updates(list(pending))
        except Exception:
            # Preserve newer updates that arrived during the failed database read.
            self._pending_vessel_updates = {**pending, **self._pending_vessel_updates}
            raise
        if not vessels:
            return
        vessels = [
            {
                **vessel,
                "action": vessel.get("action", "upsert"),
                "created": pending.get(str(vessel["mmsi"]), {}).get("_created"),
            }
            for vessel in vessels
        ]
        await manager.broadcast_vessels(
            {
                "action": "upsert",
                "source": "mixed",
                "created": None,
                "vessels": vessels,
            }
        )

    async def _load_current_vessel_updates(self, keys: list[str]) -> list[Dict[str, Any]]:
        from app.db.session import async_session_factory
        from app.services.vessel_ais_aggregation import get_current_vessels_by_mmsi

        mmsis = [int(key) for key in keys if key.isdigit()]
        vessels = []
        async with async_session_factory() as db:
            for offset in range(0, len(mmsis), VESSEL_STATE_QUERY_BATCH_SIZE):
                vessels.extend(await get_current_vessels_by_mmsi(
                    db, mmsis[offset:offset + VESSEL_STATE_QUERY_BATCH_SIZE]
                ))
        present = {int(vessel["mmsi"]) for vessel in vessels}
        vessels.extend({"mmsi": mmsi, "action": "remove"} for mmsi in mmsis if mmsi not in present)
        return vessels

    async def broadcast_vessels_periodically(self):
        while self.running:
            try:
                await self.flush_vessel_updates()
            except Exception as exc:
                logger.exception_event(
                    "Failed to flush vessel updates",
                    event="vessels.broadcast.failed",
                    context={"error": str(exc)},
                )
            await asyncio.sleep(self._vessel_flush_interval)

    async def broadcast_datasource_task_update(self, data: Dict[str, Any]):
        """Broadcast datasource task progress updates to connected clients."""
        await manager.broadcast(
            {
                "type": "data_frame",
                "channel": "datasource_tasks",
                "timestamp": to_iso8601_utc(datetime.now(UTC)),
                "payload": data,
            },
            channel="datasource_tasks",
        )

    def start(self):
        """Start all broadcasters"""
        if not self.running:
            self.running = True
            self.tasks["dashboard"] = asyncio.create_task(self.broadcast_stats(5))
            self.tasks["vessels"] = asyncio.create_task(self.broadcast_vessels_periodically())

    def stop(self):
        """Stop all broadcasters"""
        self.running = False
        for task in self.tasks.values():
            task.cancel()
        self.tasks.clear()
        self._pending_vessel_updates.clear()


broadcaster = DataBroadcaster()