import asyncio
import websockets
import json
import logging
from typing import Set

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global set to store connected clients
connected_clients: Set[websockets.WebSocketServerProtocol] = set()

async def register_client(websocket: websockets.WebSocketServerProtocol) -> None:
    """Register a new client connection."""
    connected_clients.add(websocket)
    logger.info(f"Client {websocket.remote_address} connected. Total clients: {len(connected_clients)}")
    
    # Notify all clients about new connection
    if len(connected_clients) > 1:
        notification = {
            "type": "system",
            "message": f"User from {websocket.remote_address[0]} joined the chat",
            "timestamp": asyncio.get_event_loop().time()
        }
        await broadcast_message(json.dumps(notification), exclude=websocket)

async def unregister_client(websocket: websockets.WebSocketServerProtocol) -> None:
    """Unregister a client connection."""
    if websocket in connected_clients:
        connected_clients.remove(websocket)
        logger.info(f"Client {websocket.remote_address} disconnected. Total clients: {len(connected_clients)}")
        
        # Notify remaining clients about disconnection
        if connected_clients:
            notification = {
                "type": "system",
                "message": f"User from {websocket.remote_address[0]} left the chat",
                "timestamp": asyncio.get_event_loop().time()
            }
            await broadcast_message(json.dumps(notification))

async def broadcast_message(message: str, exclude: websockets.WebSocketServerProtocol = None) -> None:
    """Broadcast a message to all connected clients except the excluded one."""
    if not connected_clients:
        return
    
    # Create a copy of the set to avoid modification during iteration
    clients_to_notify = connected_clients.copy()
    if exclude:
        clients_to_notify.discard(exclude)
    
    if clients_to_notify:
        # Send message to all clients concurrently
        await asyncio.gather(
            *[send_safe(client, message) for client in clients_to_notify],
            return_exceptions=True
        )

async def send_safe(websocket: websockets.WebSocketServerProtocol, message: str) -> None:
    """Safely send a message to a client with error handling."""
    try:
        await websocket.send(message)
    except websockets.exceptions.ConnectionClosed:
        logger.warning(f"Failed to send message to {websocket.remote_address}: connection closed")
        await unregister_client(websocket)
    except Exception as e:
        logger.error(f"Error sending message to {websocket.remote_address}: {e}")
        await unregister_client(websocket)

async def handle_client(websocket: websockets.WebSocketServerProtocol, path: str) -> None:
    """Handle a client connection."""
    await register_client(websocket)
    
    try:
        async for message in websocket:
            try:
                # Parse the incoming message
                data = json.loads(message)
                
                # Create a broadcast message with metadata
                broadcast_data = {
                    "type": "chat",
                    "message": data.get("message", ""),
                    "sender": data.get("sender", "Anonymous"),
                    "timestamp": asyncio.get_event_loop().time()
                }
                
                logger.info(f"Broadcasting message from {broadcast_data['sender']}: {broadcast_data['message']}")
                
                # Broadcast to all connected clients
                await broadcast_message(json.dumps(broadcast_data))
                
            except json.JSONDecodeError:
                logger.warning(f"Invalid JSON received from {websocket.remote_address}")
                error_response = {
                    "type": "error",
                    "message": "Invalid message format",
                    "timestamp": asyncio.get_event_loop().time()
                }
                await send_safe(websocket, json.dumps(error_response))
                
            except Exception as e:
                logger.error(f"Error processing message from {websocket.remote_address}: {e}")
                
    except websockets.exceptions.ConnectionClosed:
        logger.info(f"Connection closed by {websocket.remote_address}")
    except Exception as e:
        logger.error(f"Unexpected error with {websocket.remote_address}: {e}")
    finally:
        await unregister_client(websocket)

async def main():
    """Start the WebSocket server."""
    host = "localhost"
    port = 8765
    
    logger.info(f"Starting WebSocket server on {host}:{port}")
    
    try:
        async with websockets.serve(handle_client, host, port):
            logger.info(f"WebSocket server started successfully on ws://{host}:{port}")
            logger.info("Press Ctrl+C to stop the server")
            
            # Keep the server running
            await asyncio.Future()  # Run forever
            
    except OSError as e:
        logger.error(f"Failed to start server: {e}")
    except KeyboardInterrupt:
        logger.info("Server shutdown requested")
    except Exception as e:
        logger.error(f"Unexpected server error: {e}")

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("Server stopped")