""" Real-time WebSocket Manager for PopCorn Handles: Watch Parties, Chat, Notifications, and Live Events """ import asyncio import json import time import uuid from typing import Any from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Request, HTTPException from database import fetch_one, fetch_all, execute from auth import verify_init_data, get_current_user router = APIRouter() # ─── Connection Manager ─────────────────────────────────────────────────────── class ConnectionManager: def __init__(self): self._rooms: dict[str, set[WebSocket]] = {} # room_code -> {ws} self._user_ws: dict[int, set[WebSocket]] = {} # user_id -> {ws} self._room_owners: dict[str, int] = {} # room_code -> host_user_id self._room_media: dict[str, dict] = {} # room_code -> media info self._room_participants: dict[str, dict] = {} # room_code -> {user_id: participant_info} async def connect_to_room(self, room_code: str, ws: WebSocket, user_id: int, user_data: dict, role: str = "follower"): await ws.accept() if room_code not in self._rooms: self._rooms[room_code] = set() self._rooms[room_code].add(ws) if user_id not in self._user_ws: self._user_ws[user_id] = set() self._user_ws[user_id].add(ws) # Store participant info if room_code not in self._room_participants: self._room_participants[room_code] = {} self._room_participants[room_code][user_id] = { "user_id": user_id, "name": user_data.get("first_name", "Unknown"), "username": user_data.get("username", ""), "avatar_url": user_data.get("photo_url", ""), "role": role, "is_ready": False, "joined_at": time.time(), "is_online": True, } # Notify others await self.broadcast_to_room(room_code, { "type": "user_joined", "participant": self._room_participants[room_code][user_id], }, exclude=ws) # Send current state await ws.send_json({ "type": "room_state", "participants": list(self._room_participants[room_code].values()), "media": self._room_media.get(room_code), "room_code": room_code, "your_role": role, }) async def disconnect(self, room_code: str, ws: WebSocket, user_id: int): if room_code in self._rooms: self._rooms[room_code].discard(ws) if not self._rooms[room_code]: del self._rooms[room_code] self._room_participants.pop(room_code, None) self._room_media.pop(room_code, None) if user_id in self._user_ws: self._user_ws[user_id].discard(ws) if not self._user_ws[user_id]: del self._user_ws[user_id] # Remove participant and notify if room_code in self._room_participants: self._room_participants[room_code].pop(user_id, None) await self.broadcast_to_room(room_code, { "type": "user_left", "user_id": user_id, }) async def broadcast_to_room(self, room_code: str, message: dict, exclude: WebSocket | None = None): if room_code not in self._rooms: return dead = set() for ws in self._rooms[room_code]: if ws == exclude: continue try: await ws.send_json(message) except Exception: dead.add(ws) for ws in dead: self._rooms[room_code].discard(ws) async def send_to_user(self, user_id: int, message: dict): if user_id not in self._user_ws: return dead = set() for ws in self._user_ws[user_id]: try: await ws.send_json(message) except Exception: dead.add(ws) for ws in dead: self._user_ws[user_id].discard(ws) def get_room_participants_count(self, room_code: str) -> int: return len(self._room_participants.get(room_code, {})) manager = ConnectionManager() # ─── WebSocket Endpoint ─────────────────────────────────────────────────────── @router.websocket("/party/{room_code}") async def party_websocket(websocket: WebSocket, room_code: str): user_id = None user_data = None try: # Wait for auth message first raw = await websocket.receive_text() auth_data = json.loads(raw) if auth_data.get("type") != "auth": await websocket.close(code=4001, reason="Auth required") return init_data = auth_data.get("initData", "") user_data = verify_init_data(init_data) if not user_data or not user_data.get("id"): await websocket.close(code=4001, reason="Invalid auth") return user_id = user_data["id"] role = auth_data.get("role", "follower") # Verify room exists room = await fetch_one( "SELECT * FROM party_rooms WHERE room_code = ? AND status != 'ended'", [room_code], ) if not room: await websocket.send_json({"type": "error", "code": "ROOM_NOT_FOUND", "message": "Party room not found"}) await websocket.close(code=4004, reason="Room not found") return await manager.connect_to_room(room_code, websocket, user_id, user_data, role) # Main message loop while True: raw = await websocket.receive_text() msg = json.loads(raw) msg_type = msg.get("type") if msg_type == "ping": await websocket.send_json({"type": "pong", "timestamp": time.time()}) elif msg_type == "sync": # Leader sends sync; broadcast to all others if manager._room_participants.get(room_code, {}).get(user_id, {}).get("role") == "leader": await manager.broadcast_to_room(room_code, { "type": "sync", "is_playing": msg.get("is_playing", False), "position": msg.get("position", 0), "timestamp": time.time(), "playback_rate": msg.get("playback_rate", 1.0), }, exclude=websocket) elif msg_type == "chat": message_data = { "type": "chat_message", "message": { "id": str(uuid.uuid4()), "user_id": str(user_id), "user_name": user_data.get("first_name", "Unknown"), "content": msg.get("content", ""), "type": "chat", "reply_to_id": msg.get("reply_to_id"), "created_at": time.time(), }, } await manager.broadcast_to_room(room_code, message_data) elif msg_type == "reaction": await manager.broadcast_to_room(room_code, { "type": "reaction", "emoji": msg.get("emoji", "👍"), "target_id": msg.get("target_id"), "target_type": msg.get("target_type", "message"), "user_id": str(user_id), }) elif msg_type == "ready": is_ready = msg.get("is_ready", False) if room_code in manager._room_participants and user_id in manager._room_participants[room_code]: manager._room_participants[room_code][user_id]["is_ready"] = is_ready await manager.broadcast_to_room(room_code, { "type": "ready", "user_id": str(user_id), "is_ready": is_ready, }) elif msg_type == "buffering": await manager.broadcast_to_room(room_code, { "type": "buffering", "user_id": str(user_id), "is_buffering": msg.get("is_buffering", False), }) elif msg_type == "leave": break except WebSocketDisconnect: pass except Exception as e: print(f"[ws] error: {e}") finally: if room_code and user_id: await manager.disconnect(room_code, websocket, user_id) # ─── Chat WebSocket ──────────────────────────────────────────────────────────── @router.websocket("/chat/{conversation_id}") async def chat_websocket(websocket: WebSocket, conversation_id: str): user_id = None try: raw = await websocket.receive_text() auth_data = json.loads(raw) init_data = auth_data.get("initData", "") user_data = verify_init_data(init_data) if not user_data or not user_data.get("id"): await websocket.close(code=4001) return user_id = user_data["id"] await websocket.accept() while True: raw = await websocket.receive_text() msg = json.loads(raw) if msg.get("type") == "ping": await websocket.send_json({"type": "pong"}) elif msg.get("type") == "message": content = msg.get("content", "").strip() if content: msg_id = str(uuid.uuid4()) await execute( "INSERT INTO chat_messages (id, conversation_id, sender_id, content, created_at) VALUES (?, ?, ?, ?, datetime('now'))", [msg_id, conversation_id, str(user_id), content], ) # Update conversation last_message await execute( "UPDATE conversations SET last_message = ?, last_message_at = datetime('now') WHERE id = ?", [content, conversation_id], ) await websocket.send_json({ "type": "message_sent", "id": msg_id, "content": content, "sender_id": str(user_id), "created_at": time.time(), }) except WebSocketDisconnect: pass except Exception as e: print(f"[ws-chat] error: {e}") # ─── REST Endpoints for Room Management ─────────────────────────────────────── @router.get("/room/{room_code}/participants") async def get_room_participants(room_code: str): participants = manager._room_participants.get(room_code, {}) return {"ok": True, "data": {"participants": list(participants.values()), "count": len(participants)}} @router.post("/room/{room_code}/media") async def set_room_media(room_code: str, request: Request): body = await request.json() manager._room_media[room_code] = body await manager.broadcast_to_room(room_code, {"type": "media_changed", "media": body}) return {"ok": True} @router.post("/room/{room_code}/kick") async def kick_participant(room_code: str, request: Request): body = await request.json() target_user_id = body.get("user_id") if target_user_id: await manager.send_to_user(target_user_id, {"type": "kicked", "user_id": str(target_user_id)}) return {"ok": True} raise HTTPException(status_code=400, detail="Missing user_id")