"""WebSocket 实时通知服务：文献入库 → 推送到在线用户的 Feed"""

import asyncio

from fastapi import WebSocket
from starlette.websockets import WebSocketState


class ConnectionManager:
    """管理所有 WebSocket 连接，按用户分组"""

    def __init__(self):
        # {user_id: set[WebSocket]}
        self._connections: dict[str, set[WebSocket]] = {}
        self._lock = asyncio.Lock()

    async def connect(self, user_id: str, ws: WebSocket):
        await ws.accept()
        async with self._lock:
            self._connections.setdefault(user_id, set()).add(ws)

    async def disconnect(self, user_id: str, ws: WebSocket):
        async with self._lock:
            if user_id in self._connections:
                self._connections[user_id].discard(ws)
                if not self._connections[user_id]:
                    del self._connections[user_id]

    async def send_to_user(self, user_id: str, message: dict) -> int:
        """推送到单个用户的所有连接，返回成功数"""
        count = 0
        async with self._lock:
            connections = list(self._connections.get(user_id, set()))
        for ws in connections:
            if ws.client_state == WebSocketState.CONNECTED:
                try:
                    await ws.send_json(message)
                    count += 1
                except Exception:
                    pass
        return count

    async def broadcast(self, message: dict) -> int:
        """广播给所有在线用户"""
        count = 0
        async with self._lock:
            all_users = list(self._connections.keys())
        for uid in all_users:
            c = await self.send_to_user(uid, message)
            count += c
        return count

    async def send_new_literature_alert(self, lit_title: str, lit_pmid: int, lit_journal: str) -> int:
        """有新文献入库时，推送给所有在线用户"""
        msg = {
            "type": "new_literature",
            "title": lit_title[:100],
            "pmid": lit_pmid,
            "journal": lit_journal,
            "message": f"📖 {lit_journal}: {lit_title[:60]}...",
        }
        return await self.broadcast(msg)

    @property
    def online_count(self) -> int:
        return sum(len(v) for v in self._connections.values())

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


# 全局单例
ws_manager = ConnectionManager()
