"""
大厅聊天室 —— WebSocket 广播
"""
import json
import logging
from datetime import datetime, timezone
from fastapi import APIRouter, WebSocket, WebSocketDisconnect

logger = logging.getLogger(__name__)
router = APIRouter()


class ChatManager:
    def __init__(self):
        self.connections: dict[str, WebSocket] = {}  # user_id -> ws

    async def connect(self, user_id: str, nickname: str, ws: WebSocket):
        await ws.accept()
        self.connections[user_id] = ws
        await self.broadcast_system(f"{nickname} 进入了大厅")

    async def disconnect(self, user_id: str, nickname: str):
        self.connections.pop(user_id, None)
        await self.broadcast_system(f"{nickname} 离开了大厅")

    async def broadcast(self, message: dict):
        data = json.dumps(message, ensure_ascii=False)
        dead = []
        for uid, ws in self.connections.items():
            try:
                await ws.send_text(data)
            except Exception:
                dead.append(uid)
        for uid in dead:
            self.connections.pop(uid, None)

    async def broadcast_system(self, text: str):
        await self.broadcast({
            "type": "system",
            "content": text,
            "time": datetime.now(timezone.utc).isoformat(),
        })

    async def broadcast_chat(self, user_id: str, nickname: str, content: str):
        await self.broadcast({
            "type": "chat",
            "user_id": user_id,
            "nickname": nickname,
            "content": content,
            "time": datetime.now(timezone.utc).isoformat(),
        })

    @property
    def online_count(self):
        return len(self.connections)


chat_manager = ChatManager()


@router.websocket("/ws/chat")
async def chat_ws(ws: WebSocket):
    """
    大厅聊天 WebSocket
    客户端连接后第一条消息需发送: {"type":"auth","token":"xxx"}
    之后发送: {"type":"chat","content":"xxx"}
    """
    user_id = None
    nickname = None
    try:
        await ws.accept()
        # 第一条消息用于认证
        auth_msg = await ws.receive_text()
        auth_data = json.loads(auth_msg)
        if auth_data.get("type") != "auth":
            await ws.close(code=4001, reason="首条消息需为认证信息")
            return

        from app.auth import decode_access_token
        payload = decode_access_token(auth_data.get("token", ""))
        if not payload:
            await ws.close(code=4002, reason="Token 无效")
            return

        user_id = payload["sub"]
        nickname = payload.get("username", "Unknown")

        # 从数据库获取 nickname 太麻烦，先用 token 里的 username
        # 重新注册连接（不再调 accept，因为已经 accept 了）
        chat_manager.connections[user_id] = ws
        await chat_manager.broadcast_system(f"{nickname} 进入了大厅")

        # 发送在线人数
        await ws.send_text(json.dumps({
            "type": "online",
            "count": chat_manager.online_count,
        }))

        # 消息循环
        while True:
            raw = await ws.receive_text()
            data = json.loads(raw)
            if data.get("type") == "chat" and data.get("content", "").strip():
                await chat_manager.broadcast_chat(user_id, nickname, data["content"].strip()[:200])

    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.error(f"Chat WS error: {e}", exc_info=True)
    finally:
        if user_id:
            chat_manager.connections.pop(user_id, None)
            if nickname:
                await chat_manager.broadcast_system(f"{nickname} 离开了大厅")