"""WebSocket endpoint for real‑time progress updates."""

import json

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from backend.services.task_manager import task_manager
from backend.services.progress_publisher import register_connection, unregister_connection

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


@router.websocket("/ws/{task_id}")
async def websocket_endpoint(websocket: WebSocket, task_id: str):
    await websocket.accept()

    register_connection(task_id, websocket)

    try:
        # Send initial status
        status = task_manager.get_status(task_id)
        await websocket.send_json({
            "type": "status",
            "data": status
        })

        while True:
            data = await websocket.receive_text()
            try:
                message = json.loads(data)
                if message.get("type") == "ping":
                    await websocket.send_json({"type": "pong"})
                elif message.get("type") == "get_status":
                    status = task_manager.get_status(task_id)
                    await websocket.send_json({
                        "type": "status",
                        "data": status
                    })
            except json.JSONDecodeError:
                pass

    except WebSocketDisconnect:
        unregister_connection(task_id, websocket)
    except Exception as e:
        unregister_connection(task_id, websocket)
        print(f"WebSocket error: {e}")