# server.py (using FastAPI and aiortc)
import asyncio
import json
import uuid

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from aiortc import RTCPeerConnection, RTCSessionDescription, RTCIceCandidate

app = FastAPI()
connections = {}

@app.websocket("/ws/{room_id}/{client_id}")
async def websocket_endpoint(websocket: WebSocket, room_id: str, client_id: str):
    await websocket.accept()
    pc = RTCPeerConnection()
    connections[client_id] = pc

    @pc.on("track")
    def on_track(track):
        print("Track received")

    try:
        while True:
            message = await websocket.receive_text()
            data = json.loads(message)
            if data["type"] == "offer":
                offer = RTCSessionDescription(sdp=data["sdp"], type=data["type"])
                await pc.setRemoteDescription(offer)
                answer = await pc.createAnswer()
                await pc.setLocalDescription(answer)
                await websocket.send_text(json.dumps({"type": "answer", "sdp": pc.localDescription.sdp}))
            elif data["type"] == "answer":
                answer = RTCSessionDescription(sdp=data["sdp"], type=data["type"])
                await pc.setRemoteDescription(answer)
            elif data["type"] == "candidate":
                candidate = RTCIceCandidate(
                    component=data["candidate"]["component"],
                    foundation=data["candidate"]["foundation"],
                    ip=data["candidate"]["ip"],
                    port=data["candidate"]["port"],
                    priority=data["candidate"]["priority"],
                    protocol=data["candidate"]["protocol"],
                    type=data["candidate"]["type"],
                    sdpMid=data["candidate"]["sdpMid"],
                    sdpMLineIndex=data["candidate"]["sdpMLineIndex"],
                    usernameFragment=data["candidate"]["usernameFragment"],
                )
                await pc.addIceCandidate(candidate)
            else:
                print(f"Received unknown message: {data}")

    except WebSocketDisconnect:
        await pc.close()
        del connections[client_id]
        print(f"Client {client_id} disconnected from room {room_id}")