from typing import Dict, Set
from fastapi import WebSocket
from uuid import UUID
import json
import time
import asyncio


class ConnectionManager:
    def __init__(self):
        # room_id -> set of (websocket, user_id)
        self.active_connections: Dict[UUID, Set[tuple[WebSocket, UUID]]] = {}
        # (room_id, user_id) -> last pong timestamp
        self.last_pong: Dict[tuple[UUID, UUID], float] = {}
        # Ping interval and timeout in seconds
        self.ping_interval = 30  # Отправлять ping каждые 30 секунд
        self.ping_timeout = 60   # Отключать если нет ответа 60 секунд

    async def connect(self, websocket: WebSocket, room_id: UUID, user_id: UUID):
        await websocket.accept()
        if room_id not in self.active_connections:
            self.active_connections[room_id] = set()
        self.active_connections[room_id].add((websocket, user_id))
        # Инициализируем время последнего pong
        self.last_pong[(room_id, user_id)] = time.time()

    def disconnect(self, websocket: WebSocket, room_id: UUID, user_id: UUID):
        if room_id in self.active_connections:
            self.active_connections[room_id].discard((websocket, user_id))
            if not self.active_connections[room_id]:
                del self.active_connections[room_id]
        # Удаляем запись о последнем pong
        self.last_pong.pop((room_id, user_id), None)

    def update_pong(self, room_id: UUID, user_id: UUID):
        """Обновить время последнего pong от пользователя"""
        self.last_pong[(room_id, user_id)] = time.time()

    async def broadcast_to_room(self, room_id: UUID, message: dict, exclude_user: UUID = None):
        if room_id not in self.active_connections:
            return

        message_json = json.dumps(message, default=str)
        disconnected = []

        for websocket, user_id in self.active_connections[room_id]:
            if exclude_user and user_id == exclude_user:
                continue
            try:
                await websocket.send_text(message_json)
            except Exception:
                disconnected.append((websocket, user_id))

        for conn in disconnected:
            self.active_connections[room_id].discard(conn)

    async def send_to_user(self, room_id: UUID, user_id: UUID, message: dict):
        """Отправить сообщение конкретному пользователю в комнате"""
        if room_id not in self.active_connections:
            return False

        message_json = json.dumps(message, default=str)
        for websocket, uid in self.active_connections[room_id]:
            if uid == user_id:
                try:
                    await websocket.send_text(message_json)
                    return True
                except Exception:
                    return False
        return False

    def get_room_user_count(self, room_id: UUID) -> int:
        if room_id not in self.active_connections:
            return 0
        return len(self.active_connections[room_id])

    def get_room_users(self, room_id: UUID) -> Set[UUID]:
        """Получить список user_id в комнате"""
        if room_id not in self.active_connections:
            return set()
        return {user_id for _, user_id in self.active_connections[room_id]}

    def get_all_connections(self) -> Dict[UUID, Set[tuple[WebSocket, UUID]]]:
        """Получить все активные соединения (для фоновой задачи пинга)"""
        return self.active_connections


manager = ConnectionManager()