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)