A_Share_DP/src/ashare_dp/api/websocket/manager.py

"""WebSocket connection and subscription manager."""

from __future__ import annotations

import asyncio
import uuid

from fastapi import WebSocket
from loguru import logger


class ConnectionManager:
    """Manages WebSocket connections and per-client subscriptions.

    Tracks:
    - Active connections (client_id -> WebSocket)
    - Per-client subscriptions (client_id -> set of ts_codes)
    - Reverse index (ts_code -> set of client_ids) for fast lookup
    """

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

    async def connect(self, websocket: WebSocket) -> str:
        """Accept a new WebSocket connection and return a client_id."""
        await websocket.accept()
        client_id = str(uuid.uuid4())[:8]
        async with self._lock:
            self._connections[client_id] = websocket
            self._subscriptions[client_id] = set()
        logger.info(f"WS client connected: {client_id}")
        return client_id

    async def disconnect(self, client_id: str):
        """Remove a client and all its subscriptions."""
        async with self._lock:
            if client_id in self._subscriptions:
                # Remove from reverse index
                for code in self._subscriptions[client_id]:
                    if code in self._code_subscribers:
                        self._code_subscribers[code].discard(client_id)
                        if not self._code_subscribers[code]:
                            del self._code_subscribers[code]
                del self._subscriptions[client_id]
            self._connections.pop(client_id, None)
        logger.info(f"WS client disconnected: {client_id}")

    async def subscribe(self, client_id: str, codes: list[str]):
        """Subscribe a client to specific stock codes."""
        async with self._lock:
            if client_id not in self._subscriptions:
                return
            for code in codes:
                self._subscriptions[client_id].add(code)
                if code not in self._code_subscribers:
                    self._code_subscribers[code] = set()
                self._code_subscribers[code].add(client_id)
        logger.debug(f"Client {client_id} subscribed to {len(codes)} codes")

    async def unsubscribe(self, client_id: str, codes: list[str]):
        """Unsubscribe a client from specific codes."""
        async with self._lock:
            if client_id not in self._subscriptions:
                return
            for code in codes:
                self._subscriptions[client_id].discard(code)
                if code in self._code_subscribers:
                    self._code_subscribers[code].discard(client_id)
                    if not self._code_subscribers[code]:
                        del self._code_subscribers[code]

    async def unsubscribe_all(self, client_id: str):
        """Unsubscribe a client from all codes."""
        async with self._lock:
            if client_id not in self._subscriptions:
                return
            codes = list(self._subscriptions[client_id])
            await self.unsubscribe(client_id, codes)

    async def get_all_subscribed_codes(self) -> set[str]:
        """Get the union of all codes any client is subscribed to."""
        async with self._lock:
            return set(self._code_subscribers.keys())

    async def get_client_codes(self, client_id: str) -> set[str]:
        """Get codes a specific client is subscribed to."""
        async with self._lock:
            return self._subscriptions.get(client_id, set()).copy()

    async def send_to_client(self, client_id: str, message: dict):
        """Send a JSON message to a specific client."""
        ws = self._connections.get(client_id)
        if ws is None:
            return
        try:
            await ws.send_json(message)
        except Exception:
            await self.disconnect(client_id)

    async def broadcast(self, message: dict, client_ids: set[str] | None = None):
        """Broadcast a message to specific clients, or all if None."""
        if client_ids is None:
            client_ids = set(self._connections.keys())

        tasks = []
        for cid in client_ids:
            tasks.append(self.send_to_client(cid, message))
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)

    async def broadcast_filtered(self, spot_data: dict[str, dict]):
        """Broadcast spot data, filtering per-client based on subscriptions.

        Args:
            spot_data: Dict mapping ts_code -> spot data dict.
        """
        # Snapshot connections and subscriptions under lock
        async with self._lock:
            client_codes = {
                cid: codes.copy()
                for cid, codes in self._subscriptions.items()
                if cid in self._connections
            }

        tasks = []
        for cid, codes in client_codes.items():
            if not codes:
                continue
            # Filter to only subscribed codes for this client
            filtered = {c: spot_data[c] for c in codes if c in spot_data}
            if filtered:
                tasks.append(self.send_to_client(cid, {
                    "type": "spot",
                    "data": filtered,
                }))
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)

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

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


# Global manager instance
manager = ConnectionManager()