from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from app.realtime.voice_hub import voice_hub
from app.services.voice_rooms import voice_rooms_collection

router = APIRouter()

@router.websocket("/ws/voice")
async def voice_ws(
    ws: WebSocket,
    room_id: str = Query(...),
    username: str = Query(...)
):
    room = await voice_rooms_collection.find_one({"id": room_id})
    if not room:
        await ws.close(code=4004)
        return
    if username in voice_hub.rooms.get(room_id, {}):
        await ws.close(code=4001)
        return
    
    if len(voice_hub.rooms.get(room_id, {})) >= room.get("max_users", 5):
        await ws.close(code=4003)
        return

    await voice_hub.connect(room_id, username, ws)

    users = list(voice_hub.rooms.get(room_id, {}).keys())
    await ws.send_json({
        "type": "users",
        "users": users
    })

    await voice_hub.broadcast_except(
        room_id,
        username,
        {"type": "join", "user": username}
    )

    try:
        while True:
            msg = await ws.receive_json()
            if msg["type"] == "signal":
                await voice_hub.send_to(
                    room_id,
                    msg["to"],
                    {
                        "type": "signal",
                        "from": username,
                        "data": msg["data"]
                    }
                )
    except WebSocketDisconnect:
        voice_hub.disconnect(room_id, username)
        await voice_hub.broadcast(
            room_id,
            {"type": "leave", "user": username}
        )