"""
WebSocket Connection Manager for AI Interviews
Handles WebSocket connections, sessions, and message routing
"""
from fastapi import WebSocket, WebSocketDisconnect
from typing import Dict, Optional, Callable
from datetime import datetime
import json
import asyncio
import logging

logger = logging.getLogger(__name__)


class ConnectionManager:
    """
    Manages WebSocket connections and sessions for AI interviews
    """
    
    def __init__(self):
        # Active WebSocket connections: session_id -> WebSocket
        self.active_connections: Dict[str, WebSocket] = {}
        
        # Session metadata: session_id -> session_data
        self.sessions: Dict[str, dict] = {}
        
        # Message handlers: message_type -> handler_function
        self.message_handlers: Dict[str, Callable] = {}
        
        logger.info("WebSocket ConnectionManager initialized")
    
    async def connect(self, session_id: str, websocket: WebSocket, metadata: Optional[dict] = None):
        """
        Accept and register a new WebSocket connection
        """
        await websocket.accept()
        self.active_connections[session_id] = websocket
        
        # Initialize session data
        self.sessions[session_id] = {
            "status": "connected",
            "connected_at": datetime.utcnow().isoformat(),
            "metadata": metadata or {},
            "message_count": 0
        }
        
        logger.info(f"✅ WebSocket connected: session={session_id}")
        
        # Send connection confirmation
        await self.send_message(session_id, {
            "type": "connection",
            "status": "connected",
            "session_id": session_id,
            "timestamp": datetime.utcnow().isoformat()
        })
    
    def disconnect(self, session_id: str):
        """
        Remove a WebSocket connection and mark session as disconnected
        """
        if session_id in self.active_connections:
            del self.active_connections[session_id]
            logger.info(f"❌ WebSocket disconnected: session={session_id}")
        
        if session_id in self.sessions:
            self.sessions[session_id]["status"] = "disconnected"
            self.sessions[session_id]["disconnected_at"] = datetime.utcnow().isoformat()
    
    def is_connected(self, session_id: str) -> bool:
        """
        Check if a session is currently connected
        """
        return session_id in self.active_connections
    
    def get_session(self, session_id: str) -> Optional[dict]:
        """
        Get session metadata
        """
        return self.sessions.get(session_id)
    
    def update_session(self, session_id: str, data: dict):
        """
        Update session metadata
        """
        if session_id in self.sessions:
            self.sessions[session_id].update(data)
    
    async def send_message(self, session_id: str, message: dict):
        """
        Send a message to a specific session
        """
        if session_id in self.active_connections:
            try:
                await self.active_connections[session_id].send_json(message)
                
                # Update message count
                if session_id in self.sessions:
                    self.sessions[session_id]["message_count"] += 1
                
                logger.debug(f"📤 Message sent to session={session_id}, type={message.get('type')}")
            except Exception as e:
                logger.error(f"Error sending message to session={session_id}: {e}")
                self.disconnect(session_id)
        else:
            logger.warning(f"Cannot send message: session={session_id} not connected")
    
    async def send_text(self, session_id: str, text: str):
        """
        Send a text message to a specific session
        """
        if session_id in self.active_connections:
            try:
                await self.active_connections[session_id].send_text(text)
            except Exception as e:
                logger.error(f"Error sending text to session={session_id}: {e}")
                self.disconnect(session_id)
    
    async def send_bytes(self, session_id: str, data: bytes):
        """
        Send binary data to a specific session
        """
        if session_id in self.active_connections:
            try:
                await self.active_connections[session_id].send_bytes(data)
            except Exception as e:
                logger.error(f"Error sending bytes to session={session_id}: {e}")
                self.disconnect(session_id)
    
    async def broadcast(self, message: dict, exclude: Optional[list] = None):
        """
        Broadcast a message to all connected sessions
        """
        exclude = exclude or []
        disconnected = []
        
        for session_id, connection in self.active_connections.items():
            if session_id not in exclude:
                try:
                    await connection.send_json(message)
                except Exception as e:
                    logger.error(f"Error broadcasting to session={session_id}: {e}")
                    disconnected.append(session_id)
        
        # Clean up disconnected sessions
        for session_id in disconnected:
            self.disconnect(session_id)
    
    def register_handler(self, message_type: str, handler: Callable):
        """
        Register a message handler for a specific message type
        """
        self.message_handlers[message_type] = handler
        logger.info(f"Registered handler for message type: {message_type}")
    
    async def handle_message(self, session_id: str, message: dict):
        """
        Route message to appropriate handler
        """
        message_type = message.get("type")
        
        if message_type in self.message_handlers:
            try:
                await self.message_handlers[message_type](session_id, message)
            except Exception as e:
                logger.error(f"Error handling message type={message_type}: {e}")
                await self.send_message(session_id, {
                    "type": "error",
                    "message": f"Error processing {message_type}: {str(e)}"
                })
        else:
            logger.warning(f"No handler for message type: {message_type}")
            await self.send_message(session_id, {
                "type": "error",
                "message": f"Unknown message type: {message_type}"
            })
    
    def get_active_sessions(self) -> list:
        """
        Get list of all active session IDs
        """
        return list(self.active_connections.keys())
    
    def get_session_count(self) -> int:
        """
        Get count of active sessions
        """
        return len(self.active_connections)
    
    def cleanup_session(self, session_id: str):
        """
        Clean up session data
        """
        if session_id in self.sessions:
            del self.sessions[session_id]
            logger.info(f"🧹 Session cleaned up: session={session_id}")


# Global connection manager instance
manager = ConnectionManager()


# Utility functions
async def send_error(session_id: str, error_message: str):
    """
    Send an error message to a session
    """
    await manager.send_message(session_id, {
        "type": "error",
        "message": error_message,
        "timestamp": datetime.utcnow().isoformat()
    })


async def send_status(session_id: str, status: str, details: Optional[dict] = None):
    """
    Send a status update to a session
    """
    message = {
        "type": "status",
        "status": status,
        "timestamp": datetime.utcnow().isoformat()
    }
    
    if details:
        message.update(details)
    
    await manager.send_message(session_id, message)