"""
=========================================
CV6 AI Trading OS
WebSocket Connection Manager
=========================================
"""

import asyncio
import time
import uuid
from typing import Dict, Set
from fastapi import WebSocket
from loguru import logger


class ConnectionManager:
    """
    Manages all active WebSocket connections.
    Supports: multiple clients, symbol subscriptions,
    broadcast, heartbeat, reconnect, disconnect.
    """

    def __init__(self):
        # client_id -> WebSocket
        self.active: Dict[str, WebSocket] = {}
        # symbol -> set of client_ids
        self.subscriptions: Dict[str, Set[str]] = {}
        self._lock = asyncio.Lock()

    # ─── Connect / Disconnect ───────────────────────────────────────

    async def connect(self, ws: WebSocket, client_id: str) -> None:
        await ws.accept()
        async with self._lock:
            self.active[client_id] = ws
        logger.info(f"[WS] Connected: {client_id}  total={len(self.active)}")
        await self.send(client_id, {
            "type": "connected",
            "client_id": client_id,
            "message": "CV6 WebSocket Ready",
        })

    async def disconnect(self, client_id: str) -> None:
        async with self._lock:
            self.active.pop(client_id, None)
            for sym in list(self.subscriptions):
                self.subscriptions[sym].discard(client_id)
                if not self.subscriptions[sym]:
                    del self.subscriptions[sym]
        logger.info(f"[WS] Disconnected: {client_id}  remaining={len(self.active)}")

    # ─── Subscribe / Unsubscribe ────────────────────────────────────

    async def subscribe(self, client_id: str, symbol: str) -> None:
        symbol = symbol.upper()
        async with self._lock:
            self.subscriptions.setdefault(symbol, set()).add(client_id)
        logger.info(f"[WS] {client_id} subscribed to {symbol}")
        await self.send(client_id, {"type": "subscribed", "symbol": symbol})

    async def unsubscribe(self, client_id: str, symbol: str) -> None:
        symbol = symbol.upper()
        async with self._lock:
            if symbol in self.subscriptions:
                self.subscriptions[symbol].discard(client_id)
                if not self.subscriptions[symbol]:
                    del self.subscriptions[symbol]
        await self.send(client_id, {"type": "unsubscribed", "symbol": symbol})

    # ─── Messaging ──────────────────────────────────────────────────

    async def send(self, client_id: str, data: dict) -> None:
        ws = self.active.get(client_id)
        if ws:
            try:
                await ws.send_json(data)
            except Exception as e:
                logger.warning(f"[WS] send failed {client_id}: {e}")
                await self.disconnect(client_id)

    async def broadcast(self, data: dict) -> None:
        for cid in list(self.active):
            await self.send(cid, data)

    async def broadcast_symbol(self, symbol: str, data: dict) -> None:
        for cid in list(self.subscriptions.get(symbol.upper(), [])):
            await self.send(cid, data)

    # ─── Heartbeat ──────────────────────────────────────────────────

    async def heartbeat_loop(self, interval: int = 30) -> None:
        while True:
            await asyncio.sleep(interval)
            for cid in list(self.active):
                await self.send(cid, {"type": "heartbeat", "ts": time.time()})

    # ─── Status ─────────────────────────────────────────────────────

    def status(self) -> dict:
        return {
            "clients": len(self.active),
            "symbols": list(self.subscriptions.keys()),
            "subscriptions": {
                s: list(c) for s, c in self.subscriptions.items()
            },
        }


# Singleton
manager = ConnectionManager()