from typing import Dict, List, Optional
from fastapi import WebSocket, WebSocketDisconnect
import time
import logging

logger = logging.getLogger(__name__)

class ChatManager:
    def __init__(self):
        # 4个频道: 0, 1, 2, 3
        # 结构: {channel_id: {username: websocket}}
        self.channels: Dict[int, Dict[str, WebSocket]] = {i: {} for i in range(4)}
        # 存储聊天记录: {channel_id: [message_dict]}
        self.history: Dict[int, List[dict]] = {i: [] for i in range(4)}

    async def connect(self, websocket: WebSocket, channel_id: int, username: str):
        await websocket.accept()
        if channel_id not in self.channels:
            # 如果频道不存在，默认连入频道0或拒绝
            return 
        self.channels[channel_id][username] = websocket
        
        # 发送历史消息
        for msg in self.history[channel_id]:
            await websocket.send_json(msg)
            
        # await self.broadcast(channel_id, f"欢迎 {username} 加入频道", "system")

    def disconnect(self, channel_id: int, username: str):
        if channel_id in self.channels and username in self.channels[channel_id]:
            del self.channels[channel_id][username]
            # 离开消息可以选择不广播以减少噪音，或者仅广播给特定事件

    async def broadcast(self, channel_id: int, message: str, sender: str):
        if channel_id not in self.channels:
            return
        
        timestamp = time.strftime("%H:%M:%S", time.localtime())
        payload = {
            "type": "chat",
            "sender": sender,
            "message": message,
            "timestamp": timestamp
        }
        
        # 存储消息，限制每个频道50条
        self.history[channel_id].append(payload)
        if len(self.history[channel_id]) > 50:
            self.history[channel_id].pop(0)

        # 广播给当前所有连接用户
        for user, ws in list(self.channels[channel_id].items()):
            try:
                await ws.send_json(payload)
            except Exception as e:
                logger.error(f"Error broadcasting to {user}: {e}")
                self.disconnect(channel_id, user)