"""Data broadcaster for WebSocket connections""" import asyncio from datetime import UTC, datetime from typing import Dict, Any from app.core.time import to_iso8601_utc from app.core.websocket.manager import manager EARTH_UPDATES_CHANNEL = "earth_updates" 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 = {} vessels = [] for item in pending.values(): vessel = dict(item) source = vessel.pop("_source", None) action = vessel.pop("_action", "upsert") created = vessel.pop("_created", None) vessel["source"] = source vessel["action"] = action vessel["created"] = created vessels.append(vessel) await manager.broadcast_vessels( { "action": "upsert", "source": "mixed", "created": None, "vessels": vessels, } ) async def broadcast_vessels_periodically(self): while self.running: try: await self.flush_vessel_updates() except Exception: pass 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()