from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect

from app.auth.service import AuthService
from app.ws.manager import manager

ws_router = APIRouter()
auth_service = AuthService()


@ws_router.websocket("/ws/jobs/{job_id}")
async def websocket_endpoint(
    websocket: WebSocket,
    job_id: str,
    token: str = Query(...),
) -> None:
    """WebSocket endpoint for real-time job progress updates.

    Authentication is performed via query parameter token.
    """
    # Validate token
    claims = auth_service.validate_token(token)
    if claims is None:
        await websocket.close(code=4001, reason="Invalid or expired token")
        return

    await manager.connect(job_id, websocket)

    try:
        # Send initial connection message
        await manager.send_personal(
            websocket,
            {
                "type": "connected",
                "job_id": job_id,
                "message": "Connected to job progress stream",
            },
        )

        # Keep connection alive and listen for client messages
        while True:
            data = await websocket.receive_text()
            # Currently we don't process incoming messages,
            # but the loop keeps the connection alive
    except WebSocketDisconnect:
        manager.disconnect(job_id, websocket)
    except Exception:
        manager.disconnect(job_id, websocket)