from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from app.services.logs_service import logs_service
import asyncio

router = APIRouter()

@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """Endpoint WebSocket pour recevoir les logs en temps réel."""
    await logs_service.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            # Ici, nous pourrions ajouter la logique pour filtrer les logs
            # en fonction des préférences du client
            await logs_service.broadcast_log(data)
    except WebSocketDisconnect:
        logs_service.disconnect(websocket)

@router.websocket("/ws/{container_id}")
async def container_logs_websocket(websocket: WebSocket, container_id: str):
    """Endpoint WebSocket pour recevoir les logs d'un conteneur spécifique."""
    await logs_service.connect(websocket, container_id)
    try:
        async for log in logs_service.get_container_logs(container_id):
            await websocket.send_text(log)
    except WebSocketDisconnect:
        logs_service.disconnect(websocket, container_id)
    except Exception as e:
        await websocket.send_text(f"Error: {str(e)}")
        logs_service.disconnect(websocket, container_id)