"""
WebSocket connection handler
Presentation Layer - manages WebSocket connections and real-time communication
"""
import asyncio
from typing import List
from fastapi import WebSocket, WebSocketDisconnect
import logging

logger = logging.getLogger(__name__)

class WebSocketManager:
    """Manages WebSocket connections and broadcasting"""
    
    def __init__(self):
        self.active_connections: List[WebSocket] = []
    
    async def connect(self, websocket: WebSocket) -> None:
        """Accept and store new WebSocket connection"""
        await websocket.accept()
        self.active_connections.append(websocket)
        logger.info(f"New client connected. Total clients: {len(self.active_connections)}")
    
    def disconnect(self, websocket: WebSocket) -> None:
        """Remove WebSocket connection"""
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
            logger.info(f"Client disconnected. Total clients: {len(self.active_connections)}")
    
    async def send_personal_message(self, message: str, websocket: WebSocket) -> None:
        """Send message to specific WebSocket connection"""
        try:
            await websocket.send_text(message)
        except Exception as e:
            logger.error(f"Error sending personal message: {e}")
            self.disconnect(websocket)
    
    async def broadcast(self, message: str) -> None:
        """Broadcast message to all connected clients"""
        if not self.active_connections:
            return
        
        try:
            # Send to all connections concurrently
            tasks = [
                self._safe_send_message(connection, message) 
                for connection in self.active_connections.copy()
            ]
            
            # Execute all sends concurrently and handle exceptions
            results = await asyncio.gather(*tasks, return_exceptions=True)
            
            # Remove failed connections
            failed_connections = []
            for i, result in enumerate(results):
                if isinstance(result, Exception):
                    failed_connections.append(self.active_connections[i])
            
            for connection in failed_connections:
                self.disconnect(connection)
                
        except Exception as e:
            logger.error(f"Error in broadcast: {e}")
    
    async def _safe_send_message(self, websocket: WebSocket, message: str) -> None:
        """Safely send message to WebSocket with error handling"""
        try:
            await websocket.send_text(message)
        except WebSocketDisconnect:
            # Connection was closed
            raise
        except Exception as e:
            logger.error(f"Error sending message to client: {e}")
            raise
    
    def get_connection_count(self) -> int:
        """Get number of active connections"""
        return len(self.active_connections)
    
    async def ping_all_connections(self) -> int:
        """Ping all connections to check health, return number of healthy connections"""
        if not self.active_connections:
            return 0
        
        healthy_connections = []
        for connection in self.active_connections.copy():
            try:
                await connection.ping()
                healthy_connections.append(connection)
            except Exception:
                logger.debug("Removing unhealthy connection")
        
        self.active_connections = healthy_connections
        return len(healthy_connections)

# Global WebSocket manager instance
websocket_manager = WebSocketManager()