"""
OMNIWEB Real-time - WebSockets + Server-Sent Events
Author: Juste Elysée MALANDILA

Features:
- WebSocket connections
- Server-Sent Events (SSE)
- Broadcasting
- Rooms/Channels
- Auto-reconnect
"""

from typing import Set, Dict, List, Callable, Any, AsyncIterator
import asyncio
import json
from datetime import datetime


class WebSocketConnection:
    """WebSocket connection wrapper."""
    
    def __init__(self, websocket, client_id: str):
        self.websocket = websocket
        self.client_id = client_id
        self.rooms: Set[str] = set()
        self.metadata: Dict[str, Any] = {}
        self.connected_at = datetime.now()
    
    async def send(self, data: Any):
        """Send data to client."""
        if isinstance(data, dict):
            data = json.dumps(data)
        await self.websocket.send_text(data)
    
    async def send_json(self, data: dict):
        """Send JSON data to client."""
        await self.websocket.send_json(data)
    
    async def receive(self) -> str:
        """Receive data from client."""
        return await self.websocket.receive_text()
    
    async def receive_json(self) -> dict:
        """Receive JSON data from client."""
        return await self.websocket.receive_json()
    
    async def close(self):
        """Close connection."""
        await self.websocket.close()
    
    def join_room(self, room: str):
        """Join a room."""
        self.rooms.add(room)
    
    def leave_room(self, room: str):
        """Leave a room."""
        self.rooms.discard(room)


class RealTimeManager:
    """Manage real-time connections and broadcasting."""
    
    def __init__(self):
        self.connections: Dict[str, WebSocketConnection] = {}
        self.rooms: Dict[str, Set[str]] = {}
        self._event_handlers: Dict[str, List[Callable]] = {}
    
    def add_connection(self, client_id: str, websocket) -> WebSocketConnection:
        """Add new connection."""
        conn = WebSocketConnection(websocket, client_id)
        self.connections[client_id] = conn
        return conn
    
    def remove_connection(self, client_id: str):
        """Remove connection."""
        if client_id in self.connections:
            conn = self.connections[client_id]
            # Remove from all rooms
            for room in list(conn.rooms):
                self.leave_room(client_id, room)
            del self.connections[client_id]
    
    def get_connection(self, client_id: str) -> WebSocketConnection:
        """Get connection by client ID."""
        return self.connections.get(client_id)
    
    def join_room(self, client_id: str, room: str):
        """Add client to room."""
        if client_id in self.connections:
            self.connections[client_id].join_room(room)
            
            if room not in self.rooms:
                self.rooms[room] = set()
            self.rooms[room].add(client_id)
    
    def leave_room(self, client_id: str, room: str):
        """Remove client from room."""
        if client_id in self.connections:
            self.connections[client_id].leave_room(room)
            
            if room in self.rooms:
                self.rooms[room].discard(client_id)
                if not self.rooms[room]:
                    del self.rooms[room]
    
    async def broadcast(self, event: str, data: Any, room: str = None):
        """
        Broadcast message to all clients or specific room.
        
        Args:
            event: Event name
            data: Data to send
            room: Optional room name (broadcasts to all if None)
        """
        message = {
            "event": event,
            "data": data,
            "timestamp": datetime.now().isoformat()
        }
        
        if room:
            # Broadcast to specific room
            if room in self.rooms:
                for client_id in self.rooms[room]:
                    conn = self.connections.get(client_id)
                    if conn:
                        await conn.send_json(message)
        else:
            # Broadcast to all
            for conn in self.connections.values():
                await conn.send_json(message)
    
    async def send_to(self, client_id: str, event: str, data: Any):
        """Send message to specific client."""
        conn = self.connections.get(client_id)
        if conn:
            message = {
                "event": event,
                "data": data,
                "timestamp": datetime.now().isoformat()
            }
            await conn.send_json(message)
    
    def on(self, event: str):
        """Decorator to register event handler."""
        def decorator(func: Callable):
            if event not in self._event_handlers:
                self._event_handlers[event] = []
            self._event_handlers[event].append(func)
            return func
        return decorator
    
    async def emit_event(self, event: str, data: Any):
        """Emit event to handlers."""
        if event in self._event_handlers:
            for handler in self._event_handlers[event]:
                await handler(data)
    
    def get_stats(self) -> Dict[str, Any]:
        """Get real-time statistics."""
        return {
            "total_connections": len(self.connections),
            "total_rooms": len(self.rooms),
            "rooms": {
                room: len(clients) 
                for room, clients in self.rooms.items()
            }
        }


# Global manager instance
_realtime_manager = RealTimeManager()


def get_realtime_manager() -> RealTimeManager:
    """Get global real-time manager."""
    return _realtime_manager


def enable_realtime(app):
    """
    Enable real-time features on OMNIWEB app.
    
    Example:
        >>> from omniweb import OmniWeb
        >>> from omniweb.realtime import enable_realtime, broadcast
        >>> 
        >>> app = OmniWeb()
        >>> enable_realtime(app)
        >>> 
        >>> @app.websocket("/ws")
        >>> async def websocket_endpoint(websocket):
        >>>     await websocket.accept()
        >>>     
        >>>     # Auto-managed connection
        >>>     client_id = "user-123"
        >>>     conn = app.realtime.add_connection(client_id, websocket)
        >>>     
        >>>     try:
        >>>         while True:
        >>>             data = await conn.receive_json()
        >>>             # Echo back
        >>>             await broadcast("message", data)
        >>>     finally:
        >>>         app.realtime.remove_connection(client_id)
    """
    app.realtime = get_realtime_manager()
    return app.realtime


async def broadcast(event: str, data: Any, room: str = None):
    """
    Broadcast message to all connected clients.
    
    Example:
        >>> from omniweb.realtime import broadcast
        >>> 
        >>> @app.post("/notify")
        >>> async def notify_all(message: str):
        >>>     await broadcast("notification", {"message": message})
        >>>     return {"status": "broadcasted"}
    """
    manager = get_realtime_manager()
    await manager.broadcast(event, data, room)


class ServerSentEvents:
    """
    Server-Sent Events (SSE) for real-time updates.
    
    Example:
        >>> from omniweb.realtime import ServerSentEvents
        >>> 
        >>> @app.get("/stream")
        >>> async def stream_updates():
        >>>     async def event_generator():
        >>>         for i in range(10):
        >>>             yield {"data": f"Update {i}"}
        >>>             await asyncio.sleep(1)
        >>>     
        >>>     return ServerSentEvents(event_generator())
    """
    
    def __init__(self, event_source: AsyncIterator[dict]):
        self.event_source = event_source
    
    async def __call__(self, scope, receive, send):
        """ASGI callable."""
        await send({
            'type': 'http.response.start',
            'status': 200,
            'headers': [
                (b'content-type', b'text/event-stream'),
                (b'cache-control', b'no-cache'),
                (b'connection', b'keep-alive'),
            ],
        })
        
        async for event in self.event_source:
            data = json.dumps(event.get('data', {}))
            message = f"data: {data}\n\n"
            
            await send({
                'type': 'http.response.body',
                'body': message.encode('utf-8'),
                'more_body': True,
            })
        
        await send({
            'type': 'http.response.body',
            'body': b'',
            'more_body': False,
        })


def realtime_route(room: str = None):
    """
    Decorator for creating real-time routes.
    
    Example:
        >>> from omniweb.realtime import realtime_route
        >>> 
        >>> @app.get("/live")
        >>> @realtime_route(room="updates")
        >>> async def live_updates():
        >>>     async def generator():
        >>>         while True:
        >>>             data = await get_latest_update()
        >>>             yield data
        >>>             await asyncio.sleep(1)
        >>>     
        >>>     return generator()
    """
    def decorator(func: Callable):
        async def wrapper(*args, **kwargs):
            result = await func(*args, **kwargs)
            
            # If result is an async generator, wrap in SSE
            if hasattr(result, '__aiter__'):
                return ServerSentEvents(result)
            
            return result
        
        return wrapper
    return decorator
