"""
WebSocket进度路由
"""
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from app.core.progress_manager import progress_manager
import logging

logger = logging.getLogger(__name__)

router = APIRouter()


@router.websocket("/progress/{task_id}")
async def websocket_progress(
    websocket: WebSocket,
    task_id: str
):
    """WebSocket进度推送"""
    await progress_manager.connect(websocket, task_id)
    
    try:
        while True:
            data = await websocket.receive_text()
            if data == "ping":
                await websocket.send_text("pong")
            elif data == "close":
                break
    except WebSocketDisconnect:
        logger.info(f"WebSocket断开连接: task_id={task_id}")
    except Exception as e:
        logger.error(f"WebSocket错误: {e}")
    finally:
        await progress_manager.disconnect(websocket, task_id)