from fastapi import WebSocket
from typing import List
from app.utils.logger import get_logger

logger = get_logger(__name__)


class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
        logger.info(f"WebSocket connected. Total connections: {len(self.active_connections)}")

    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
            logger.info(f"WebSocket disconnected. Total connections: {len(self.active_connections)}")

    async def broadcast(self, event_type: str, data: dict):
        message = {"event_type": event_type, "data": data, "timestamp": None}
        disconnected = []

        for connection in self.active_connections:
            try:
                await connection.send_json(message)
            except Exception:
                disconnected.append(connection)

        for conn in disconnected:
            self.disconnect(conn)

    async def broadcast_price_update(self, region: str, price_data: dict):
        await self.broadcast("price_update", {"region": region, "price_data": price_data})

    async def broadcast_battery_update(self, battery_id: str, battery_state: dict):
        await self.broadcast("battery_update", {"battery_id": battery_id, "battery_state": battery_state})

    async def broadcast_trade(self, trade: dict):
        await self.broadcast("trade_executed", trade)

    async def broadcast_alert(self, alert: dict):
        await self.broadcast("alert_triggered", alert)

    async def broadcast_backtest_progress(self, backtest_id: str, progress: float, status: str):
        await self.broadcast(
            "backtest_progress",
            {"backtest_id": backtest_id, "progress": progress, "status": status},
        )

    async def broadcast_model_training_progress(
        self, model_id: str, progress: float, epoch: int = None, metrics: dict = None
    ):
        await self.broadcast(
            "model_training_progress",
            {"model_id": model_id, "progress": progress, "epoch": epoch, "metrics": metrics},
        )


manager = ConnectionManager()