"""WebSocket 实时推送"""
import json

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

from app.services.dashboard_stats import get_dashboard_stats

router = APIRouter(prefix="/api", tags=["websocket"])


@router.websocket("/ws")
async def websocket_dashboard(websocket: WebSocket):
    app = websocket.app
    await websocket.accept()
    manager = app.state.ws_manager
    await manager.connect(websocket)
    try:
        stats = await get_dashboard_stats(app.state.scheduler.running)
        await websocket.send_json({"type": "stats", "data": stats})
        while True:
            raw = await websocket.receive_text()
            try:
                msg = json.loads(raw)
            except json.JSONDecodeError:
                continue
            if msg.get("type") == "ping":
                await websocket.send_json({"type": "pong"})
    except WebSocketDisconnect:
        pass
    finally:
        await manager.disconnect(websocket)