129 lines
4.8 KiB
Python
129 lines
4.8 KiB
Python
from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect, HTTPException
|
|
from sqlalchemy.orm import Session
|
|
from typing import Dict, List
|
|
from pydantic import BaseModel
|
|
import random
|
|
import string
|
|
from app.database import get_db
|
|
from app.models.models import Room, User, Song, Playlist
|
|
from app.auth import get_current_user
|
|
|
|
router = APIRouter()
|
|
|
|
class RoomCreate(BaseModel):
|
|
name: str
|
|
|
|
class ConnectionManager:
|
|
def __init__(self):
|
|
self.active_connections: Dict[str, List[dict]] = {}
|
|
|
|
async def connect(self, websocket: WebSocket, room_code: str, username: str):
|
|
await websocket.accept()
|
|
if room_code not in self.active_connections:
|
|
self.active_connections[room_code] = []
|
|
|
|
# Проверяем, не подключен ли уже этот пользователь
|
|
existing = [conn for conn in self.active_connections[room_code] if conn["username"] == username]
|
|
if not existing:
|
|
self.active_connections[room_code].append({"ws": websocket, "username": username})
|
|
|
|
await self.broadcast({"type": "user_count", "count": len(self.active_connections[room_code])}, room_code)
|
|
|
|
def disconnect(self, websocket: WebSocket, room_code: str):
|
|
if room_code in self.active_connections:
|
|
self.active_connections[room_code] = [
|
|
conn for conn in self.active_connections[room_code] if conn["ws"] != websocket
|
|
]
|
|
return len(self.active_connections[room_code])
|
|
return 0
|
|
|
|
async def broadcast(self, message: dict, room_code: str):
|
|
if room_code in self.active_connections:
|
|
disconnected = []
|
|
for connection in self.active_connections[room_code]:
|
|
try:
|
|
await connection["ws"].send_json(message)
|
|
except:
|
|
disconnected.append(connection)
|
|
|
|
# Удаляем отключенные соединения
|
|
for conn in disconnected:
|
|
if conn in self.active_connections[room_code]:
|
|
self.active_connections[room_code].remove(conn)
|
|
|
|
def get_user_count(self, room_code: str):
|
|
if room_code in self.active_connections:
|
|
return len(self.active_connections[room_code])
|
|
return 0
|
|
|
|
manager = ConnectionManager()
|
|
|
|
def generate_room_code():
|
|
return ''.join(random.choices(string.ascii_uppercase + string.digits, k=6))
|
|
|
|
@router.post("/create")
|
|
def create_room(
|
|
room: RoomCreate,
|
|
current_user: User = Depends(get_current_user),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
code = generate_room_code()
|
|
db_room = Room(name=room.name, code=code, creator_id=current_user.id)
|
|
db.add(db_room)
|
|
db.commit()
|
|
db.refresh(db_room)
|
|
return db_room
|
|
|
|
@router.get("/{room_code}")
|
|
def get_room(room_code: str, db: Session = Depends(get_db)):
|
|
room = db.query(Room).filter(Room.code == room_code).first()
|
|
if not room:
|
|
raise HTTPException(status_code=404, detail="Room not found")
|
|
user_count = manager.get_user_count(room_code)
|
|
return {**room.__dict__, "user_count": user_count}
|
|
|
|
@router.post("/{room_code}/add-song/{song_id}")
|
|
def add_song_to_room(
|
|
room_code: str,
|
|
song_id: int,
|
|
current_user: User = Depends(get_current_user),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
room = db.query(Room).filter(Room.code == room_code).first()
|
|
if not room:
|
|
raise HTTPException(status_code=404, detail="Room not found")
|
|
|
|
song = db.query(Song).filter(Song.id == song_id).first()
|
|
if not song:
|
|
raise HTTPException(status_code=404, detail="Song not found")
|
|
|
|
return {"message": "Song added to queue"}
|
|
|
|
@router.post("/{room_code}/add-playlist/{playlist_id}")
|
|
def add_playlist_to_room(
|
|
room_code: str,
|
|
playlist_id: int,
|
|
current_user: User = Depends(get_current_user),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
room = db.query(Room).filter(Room.code == room_code).first()
|
|
if not room:
|
|
raise HTTPException(status_code=404, detail="Room not found")
|
|
|
|
playlist = db.query(Playlist).filter(Playlist.id == playlist_id).first()
|
|
if not playlist:
|
|
raise HTTPException(status_code=404, detail="Playlist not found")
|
|
|
|
return {"message": "Playlist added to queue", "songs": [{"id": s.id, "title": s.title, "artist": s.artist} for s in playlist.songs]}
|
|
|
|
@router.websocket("/ws/{room_code}")
|
|
async def websocket_endpoint(websocket: WebSocket, room_code: str, username: str = "User"):
|
|
await manager.connect(websocket, room_code, username)
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_json()
|
|
await manager.broadcast(data, room_code)
|
|
except WebSocketDisconnect:
|
|
count = manager.disconnect(websocket, room_code)
|
|
await manager.broadcast({"type": "user_count", "count": count}, room_code)
|