"""WebSocket endpoint for live event streaming."""

from __future__ import annotations

import asyncio

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from mixersystem.studio.event_bus import StudioEvent, EventBus

router = APIRouter()

_HEARTBEAT_INTERVAL = 30


@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()

    event_bus: EventBus = websocket.app.state.event_bus
    queue = event_bus.subscribe()

    try:
        while True:
            try:
                event: StudioEvent = await asyncio.wait_for(
                    queue.get(), timeout=_HEARTBEAT_INTERVAL,
                )
                await websocket.send_text(event.to_json())
            except asyncio.TimeoutError:
                # Send heartbeat
                heartbeat = StudioEvent(type="heartbeat")
                await websocket.send_text(heartbeat.to_json())
    except WebSocketDisconnect:
        pass
    except (asyncio.CancelledError, ConnectionError, RuntimeError):
        pass
    finally:
        event_bus.unsubscribe(queue)
