"""
WebSocket连接管理模块
负责管理WebSocket连接和广播消息
"""
import asyncio
import logging
from typing import List
from fastapi import WebSocket

logger = logging.getLogger(__name__)

class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []
        self._lock = asyncio.Lock()  # 添加锁以保护连接列表

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        async with self._lock:
            # 检查连接是否已经存在，如果存在则不重复添加
            if websocket not in self.active_connections:
                self.active_connections.append(websocket)
                logger.info(f"WebSocket客户端连接，当前连接数: {len(self.active_connections)}")

    async def disconnect(self, websocket: WebSocket):
        async with self._lock:
            # 安全移除连接，如果不存在则忽略
            if websocket in self.active_connections:
                self.active_connections.remove(websocket)
                logger.info(f"WebSocket客户端断开，当前连接数: {len(self.active_connections)}")

    async def broadcast(self, data):
        """向所有连接的客户端广播数据"""
        async with self._lock:
            # 复制列表以避免在迭代时修改
            connections = self.active_connections.copy()
        
        # 在锁外处理发送操作
        disconnect_list = []
        for connection in connections:
            try:
                await connection.send_json(data)
            except Exception as e:
                logger.error(f"发送消息失败: {e}")
                disconnect_list.append(connection)
        
        # 移除断开的连接
        if disconnect_list:
            async with self._lock:
                for conn in disconnect_list:
                    if conn in self.active_connections:
                        self.active_connections.remove(conn)
                        logger.info(f"移除断开的连接，剩余连接数: {len(self.active_connections)}")