"""WebSocket /ws?token=<jwt> — real-time event stream"""
from __future__ import annotations

import asyncio
import logging

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from web.auth import verify_ws_token
from web.ws_bus import get_bus

router = APIRouter()
logger = logging.getLogger(__name__)


@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket, token: str = ""):
    """
    Authenticate via JWT query param or ttb_session cookie, then stream events from WSBus.

    Events common to all users: status_snapshot, queue_update, node_executing, server_state
    Events private to submitter: generation_complete, generation_error
    """
    payload = verify_ws_token(token)
    if payload is None:
        # Fallback: browsers send cookies automatically with WebSocket connections
        cookie_token = websocket.cookies.get("ttb_session", "")
        payload = verify_ws_token(cookie_token)
    if payload is None:
        await websocket.close(code=4001, reason="Unauthorized")
        return

    user_label: str = payload.get("sub", "anonymous")
    bus = get_bus()
    queue = bus.subscribe(user_label)
    await websocket.accept()
    logger.info("WS connected: user=%s", user_label)

    try:
        while True:
            # Wait for an event from the bus
            try:
                frame = await asyncio.wait_for(queue.get(), timeout=30.0)
            except asyncio.TimeoutError:
                # Send a keepalive ping
                try:
                    await websocket.send_text('{"type":"ping"}')
                except Exception:
                    break
                continue

            try:
                await websocket.send_text(frame)
            except Exception:
                break

    except WebSocketDisconnect:
        pass
    finally:
        bus.unsubscribe(user_label, queue)
        logger.info("WS disconnected: user=%s", user_label)