"""WebSocket endpoint for real-time version room events."""

import uuid

from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect

from rehearsalhub.repositories.member import MemberRepository
from rehearsalhub.db.engine import get_session
from rehearsalhub.services.auth import decode_token
from rehearsalhub.ws import manager

router = APIRouter(tags=["websocket"])


@router.websocket("/ws/versions/{version_id}")
async def version_ws(
    version_id: uuid.UUID,
    websocket: WebSocket,
    token: str = Query(...),
):
    """WebSocket endpoint. Requires a valid JWT passed as ?token=<jwt>."""
    # Validate token before accepting the connection
    async for session in get_session():
        try:
            payload = decode_token(token)
            member_id = uuid.UUID(payload["sub"])
            member = await MemberRepository(session).get_by_id(member_id)
            if member is None:
                raise ValueError("member not found")
        except Exception:
            await websocket.close(code=4001)
            return

    await manager.connect(version_id, websocket)
    try:
        while True:
            # Echo back any client pings; clients can send {"event": "ping"}
            data = await websocket.receive_json()
            if data.get("event") == "ping":
                await websocket.send_json({"event": "pong"})
    except WebSocketDisconnect:
        manager.disconnect(version_id, websocket)