from typing import Dict, List
from fastapi import WebSocket
from app.core.json_helpers import _dumps
import logging

logger = logging.getLogger(__name__)


class WebSocketManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
    
    async def connect(self, websocket: WebSocket, client_id: str):
        """클라이언트 연결"""
        await websocket.accept()
        self.active_connections[client_id] = websocket
        logger.info(f"Client {client_id} connected. Total connections: {len(self.active_connections)}")
    
    def disconnect(self, client_id: str):
        """클라이언트 연결 해제"""
        if client_id in self.active_connections:
            del self.active_connections[client_id]
            logger.info(f"Client {client_id} disconnected. Total connections: {len(self.active_connections)}")
    
    async def send_personal_message(self, message: str, client_id: str):
        """특정 클라이언트에게 메시지 전송"""
        if client_id in self.active_connections:
            try:
                await self.active_connections[client_id].send_text(message)
            except Exception as e:
                logger.error(f"Error sending message to {client_id}: {e}")
                self.disconnect(client_id)
    
    async def broadcast(self, message: str):
        """모든 연결된 클라이언트에게 메시지 브로드캐스트"""
        disconnected_clients = []
        for client_id, connection in self.active_connections.items():
            try:
                await connection.send_text(message)
            except Exception as e:
                logger.error(f"Error broadcasting to {client_id}: {e}")
                disconnected_clients.append(client_id)
        
        # 연결 실패한 클라이언트 제거
        for client_id in disconnected_clients:
            self.disconnect(client_id)
    
    async def broadcast_json(self, data: dict):
        """JSON 데이터를 모든 클라이언트에게 브로드캐스트"""
        message = _dumps(data)
        await self.broadcast(message)
    
    async def send_market_update(self, stock_code: str, data: dict):
        """시장 데이터 업데이트 전송"""
        message = {
            "type": "market_update",
            "stock_code": stock_code,
            "data": data
        }
        await self.broadcast_json(message)
    
    async def send_trade_notification(self, trade_data: dict):
        """거래 알림 전송"""
        message = {
            "type": "trade_notification",
            "data": trade_data
        }
        await self.broadcast_json(message)
    
    async def send_mcp_status(self, status: dict):
        """MCP 상태 업데이트 전송"""
        message = {
            "type": "mcp_status",
            "data": status
        }
        await self.broadcast_json(message)


websocket_manager = WebSocketManager()