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.time import to_iso8601_utc
from app.core.websocket.manager import manager



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,
        )

    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="all",
        )

    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()