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

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

logger = logging.getLogger("pingwatch.ws")

router = APIRouter()


class ConnectionManager:
    """WebSocket 连接管理器"""

    def __init__(self):
        self._connections: set[WebSocket] = set()

    async def connect(self, ws: WebSocket):
        await ws.accept()
        self._connections.add(ws)

    def disconnect(self, ws: WebSocket):
        self._connections.discard(ws)

    async def broadcast(self, message: dict):
        """向所有客户端广播消息"""
        dead = set()
        for ws in self._connections:
            try:
                await ws.send_json(message)
            except Exception:
                dead.add(ws)
        self._connections -= dead

    @property
    def count(self) -> int:
        return len(self._connections)


manager = ConnectionManager()


@router.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):
    """WebSocket 端点，用于前端实时接收状态更新"""
    await manager.connect(ws)
    try:
        while True:
            # 保持连接，接收心跳 pong
            data = await ws.receive_text()
            if data == "ping":
                await ws.send_text("pong")
    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.error(f"WebSocket 异常: {e}")
    finally:
        manager.disconnect(ws)


async def broadcast_status_change(device_id: int, status: str, name: str):
    """广播设备状态变化"""
    await manager.broadcast({
        "type": "device_status_change",
        "device_id": device_id,
        "status": status,
        "name": name,
        "timestamp": asyncio.get_event_loop().time(),
    })