from typing import List
from enum import Enum
from fastapi import APIRouter, WebSocket
from starlette.websockets import WebSocketDisconnect
import logging

url = '/ws'
swag_tags:List[str|Enum] = ["App Socket"]
router = APIRouter(tags=swag_tags, prefix=url)

logger = logging.getLogger(__name__)

def _get_active_websockets(websocket: WebSocket) -> set[WebSocket]:
    if not hasattr(websocket.app.state, "active_websockets"):
        websocket.app.state.active_websockets = set[WebSocket]()
    return websocket.app.state.active_websockets

@router.websocket('')
@router.websocket('/')
async def websocket_endpoint(websocket: WebSocket) -> None:
    await websocket.accept()
    active_websockets = _get_active_websockets(websocket)
    active_websockets.add(websocket)

    while True:
        try:
            request_processor = getattr(websocket.app.state, "request_processor", None)
            if request_processor is None:
                logger.error("request_processor indisponible dans app.state")
                break
            data = await websocket.receive_json()
            response = await request_processor.process_request(data)
            if response:
                try:
                    if isinstance(response, dict):
                        await websocket.send_json(response)
                    else:
                        for item in response:
                            await websocket.send_json(item) 
                except Exception as e:
                    logger.error(f"Error sending response to WebSocket: {e}")
        except WebSocketDisconnect:
            break
        except Exception as e:
            logger.error(f"WebSocket error: {e}")
            break
    active_websockets.discard(websocket)
    await websocket.close()