"""WebSocket Connection Manager"""
import json
from datetime import datetime
from typing import Set, Dict, List
from fastapi import WebSocket


class WebSocketManager:
    def __init__(self):
        self.active_connections: Set[WebSocket] = set()
        self.subscriptions: Dict[WebSocket, Dict] = {}

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.add(websocket)
        self.subscriptions[websocket] = {"asset": None, "timeframe": None}

    def disconnect(self, websocket: WebSocket):
        self.active_connections.discard(websocket)
        self.subscriptions.pop(websocket, None)

    async def subscribe(self, websocket: WebSocket, asset: str = None, timeframe: str = None):
        self.subscriptions[websocket] = {"asset": asset, "timeframe": timeframe}

    async def broadcast_signal(self, signal: dict):
        disconnected = set()
        for ws in self.active_connections:
            try:
                sub = self.subscriptions.get(ws, {})
                asset_match = sub.get("asset") is None or sub["asset"] == signal["asset"]
                tf_match = sub.get("timeframe") is None or sub["timeframe"] == signal["timeframe"]
                if asset_match and tf_match:
                    await ws.send_json({"type": "signal", "data": signal})
            except Exception:
                disconnected.add(ws)
        for ws in disconnected:
            self.disconnect(ws)

    async def broadcast_all(self, message: dict):
        disconnected = set()
        for ws in self.active_connections:
            try:
                await ws.send_json(message)
            except Exception:
                disconnected.add(ws)
        for ws in disconnected:
            self.disconnect(ws)