"""
WebSocket endpoint — one persistent connection per connected client (waiter PWA, KDS).

Authentication: JWT token passed as query param ?token=<jwt>
(Same pattern as SSE — browser WebSocket API also cannot set custom headers.)

Protocol:
  1. Client connects: ws://host/api/ws?token=<jwt>
  2. Client immediately sends: { "cursor": <last_seq_id> }
     (0 = first connect / no history; any positive int = resume from that point)
  3. Server replays all missed events (seq > cursor) as individual frames
  4. Server sends { "type": "ready" } to signal end of replay / start of live stream
  5. Live events arrive as: { "seq": 123, "type": "...", "data": { ... } }
  6. Every 25s server sends: { "type": "ping" }
  7. Client may respond: { "type": "pong" } (ignored if not sent)
  8. On disconnect/error the client reconnects and repeats from step 2
"""

import asyncio
import json
import logging

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from starlette.websockets import WebSocketState

from routers.deps import decode_token
from services.ws_bus import connect, disconnect, get_events_since

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

KEEPALIVE_INTERVAL = 25  # seconds


@router.websocket("/connect")
async def ws_endpoint(websocket: WebSocket, token: str = Query(...)):
    # ── Auth ──────────────────────────────────────────────────────────────────
    try:
        payload = decode_token(token)
        user_id: int = int(payload["sub"])
    except Exception:
        await websocket.close(code=4001)
        return

    await websocket.accept()

    # ── Wait for cursor frame ─────────────────────────────────────────────────
    cursor = 0
    try:
        raw = await asyncio.wait_for(websocket.receive_text(), timeout=10.0)
        msg = json.loads(raw)
        cursor = int(msg.get("cursor", 0))
    except (asyncio.TimeoutError, Exception):
        pass  # no cursor sent — treat as fresh connect (cursor=0)

    # ── Register this connection ──────────────────────────────────────────────
    q = await connect(user_id)

    try:
        # ── Replay missed events ──────────────────────────────────────────────
        if cursor > 0:
            missed = get_events_since(cursor, user_id)
            for event in missed:
                await websocket.send_text(json.dumps(event))

        # Signal end of replay / start of live stream
        await websocket.send_text(json.dumps({"type": "ready"}))

        # ── Live stream loop ──────────────────────────────────────────────────
        while True:
            try:
                frame = await asyncio.wait_for(q.get(), timeout=KEEPALIVE_INTERVAL)
                if websocket.client_state == WebSocketState.CONNECTED:
                    await websocket.send_text(frame)
            except asyncio.TimeoutError:
                # Send keepalive ping
                if websocket.client_state == WebSocketState.CONNECTED:
                    await websocket.send_text(json.dumps({"type": "ping"}))
                else:
                    break

    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.debug("ws_endpoint: connection closed for user %s: %s", user_id, e)
    finally:
        await disconnect(user_id, q)