"""WebSocket handler for streaming model updates.

Provides FastAPI WebSocket endpoints for real-time event streaming.

⚠️ AUTO-GENERATED BY PRISM - DO NOT EDIT
"""

from __future__ import annotations

import asyncio
import json
import logging

from fastapi import WebSocket, WebSocketDisconnect

from ._generated.event_bus import StreamEvent, event_bus

logger = logging.getLogger(__name__)


async def handle_websocket(
    websocket: WebSocket,
    *,
    model: str | None = None,
    event_types: list[str] | None = None,
    heartbeat_interval: int = {{ heartbeat_interval }},
) -> None:
    """Handle a WebSocket connection for streaming events.

    Accepts the connection, then forwards matching events from the
    event bus as JSON messages. Sends periodic heartbeat pings.

    Args:
        websocket: The WebSocket connection.
        model: Filter to events for this model name only.
        event_types: Filter to these event types only.
        heartbeat_interval: Seconds between heartbeat pings.
    """
    await websocket.accept()
    logger.info("WebSocket connected (model=%s, events=%s)", model, event_types)

    subscriber = event_bus.subscribe(model=model, event_types=event_types)

    async def _send_events() -> None:
        try:
            async for event in subscriber:
                data = {
                    "model": event.model,
                    "event_type": event.event_type,
                    "entity_id": event.entity_id,
                }
                if event.payload is not None:
                    data["payload"] = event.payload
                if event.metadata:
                    data["metadata"] = event.metadata
                await websocket.send_text(json.dumps(data, default=str))
        except (WebSocketDisconnect, asyncio.CancelledError):
            pass

    async def _heartbeat() -> None:
        try:
            while True:
                await asyncio.sleep(heartbeat_interval)
                await websocket.send_text(json.dumps({"type": "heartbeat"}))
        except (WebSocketDisconnect, asyncio.CancelledError):
            pass

    # Run event sending and heartbeat concurrently
    send_task = asyncio.create_task(_send_events())
    heartbeat_task = asyncio.create_task(_heartbeat())

    try:
        # Wait for client messages (to detect disconnection)
        while True:
            data = await websocket.receive_text()
            # Client can send filter updates
            try:
                msg = json.loads(data)
                if msg.get("type") == "ping":
                    await websocket.send_text(json.dumps({"type": "pong"}))
            except json.JSONDecodeError:
                pass
    except WebSocketDisconnect:
        logger.info("WebSocket disconnected (model=%s)", model)
    finally:
        send_task.cancel()
        heartbeat_task.cancel()
        await asyncio.gather(send_task, heartbeat_task, return_exceptions=True)


__all__ = ["handle_websocket"]
