main.py

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Depends
from jose import jwt, JWTError
from typing import Dict, List
import uuid

app = FastAPI()

SECRET_KEY = "secret"
ALGORITHM = "HS256"

sessions: Dict[str, dict] = {}

def verify_token(token: str):
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        return payload
    except JWTError:
        raise HTTPException(status_code=401, detail="Invalid Token")

@app.post("/create-session")
def create_session(token: str):
    user = verify_token(token)
    session_code = str(uuid.uuid4())[:6]

    sessions[session_code] = {
        "host_id": user["id"],
        "members": [],
        "connections": {}
    }

    return {"session_code": session_code}

@app.websocket("/ws/{session_code}")
async def websocket_endpoint(websocket: WebSocket, session_code: str, token: str):
    user = verify_token(token)
    await websocket.accept()

    if session_code not in sessions:
        await websocket.close()
        return

    session = sessions[session_code]

    session["members"].append(user["id"])
    session["connections"][user["id"]] = websocket

    for connection in session["connections"].values():
        await connection.send_json({
            "event": "member-joined",
            "user_id": user["id"]
        })

    try:
        while True:
            data = await websocket.receive_json()

            for connection in session["connections"].values():
                await connection.send_json({
                    "event": "receive-message",
                    "sender": user["id"],
                    "message": data["message"]
                })

    except WebSocketDisconnect:

        session["members"].remove(user["id"])
        del session["connections"][user["id"]]

        if user["id"] == session["host_id"]:
            for connection in session["connections"].values():
                await connection.send_json({
                    "event": "session-ended"
                })
            del sessions[session_code]
            return

        for connection in session["connections"].values():
            await connection.send_json({
                "event": "member-left",
                "user_id": user["id"]
            })