Source code for api.v1.websocket

import asyncio
import logging
from typing import Set

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends
from starlette.websockets import WebSocketState
from utils.i18n import get_translations, _
from babel.support import Translations

from .network import update_container_position, get_network_data

router = APIRouter()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)



class WebSocketManager:
    """Класс для управления подключениями к WebSocket-серверу.

    Обрабатывает подключение, отключение и трансляцию сообщений для клиентов,
    подключённых через WebSocket.
    """

    def __init__(self):
        """Инициализирует менеджер подключений к WebSocket-серверу.

        Создаёт пустое множество для хранения активных соединений.
        """
        self.active_connections: Set[WebSocket] = set()



    async def connect(self, websocket: WebSocket):
        """Подключает клиента к WebSocket-серверу.

        Принимает подключение и добавляет WebSocket в список активных соединений.

        Args:
            websocket (WebSocket): Объект WebSocket для подключения.

        Raises:
            Exception: Если произошла ошибка при принятии подключения.
        """
        await websocket.accept()
        self.active_connections.add(websocket)




    def disconnect(self, websocket: WebSocket):
        """Отключает клиента от WebSocket-серверу.

        Удаляет WebSocket из списка активных соединений.

        Args:
            websocket (WebSocket): Объект WebSocket для отключения.
        """
        self.active_connections.remove(websocket)




    async def broadcast(self, message: dict):
        """Транслирует сообщения всем подключённым клиентам.

        Отправляет JSON-данные всем активным соединениям, обрабатывая
        отключения клиентов.

        Args:
            message (dict): Сообщение для трансляции в формате JSON.

        Raises:
            WebSocketDisconnect: Если клиент неожиданно отключился.
            Exception: Если произошла другая ошибка при отправке.
        """
        for connection in list(self.active_connections):
            try:
                if connection.client_state == WebSocketState.CONNECTED:
                    await connection.send_json(message)
            except WebSocketDisconnect:
                self.disconnect(connection)
            except Exception as e:
                logger.error(f"Broadcast error: {e}")
                self.disconnect(connection)



manager = WebSocketManager()



@router.websocket("/api/v1")
async def websocket_endpoint(websocket: WebSocket, translations: Translations = Depends(get_translations())):
    """Обработчик WebSocket-соединения.

    Управляет соединением клиента, принимает данные для обновления позиций
    контейнеров и транслирует сетевые данные.

    Args:
        websocket (WebSocket): Объект WebSocket для обработки соединения.
        translations: Объект переводов для текущего языка.

    Raises:
        WebSocketDisconnect: Если клиент отключился.
        Exception: Если произошла ошибка при обработке данных.
    """
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_json()
            logger.info(f"Received data: {data}")

            if all(key in data for key in ["container_id", "container_x", "container_y"]):
                await update_container_position(data)

            network_data = await get_network_data()
            logger.info(f"Current network data: {network_data}")

            await manager.broadcast(network_data)
            await asyncio.sleep(2)
    except WebSocketDisconnect:
        logger.info("Client disconnected")
    except Exception as e:
        logger.error(f"Error in websocket_endpoint: {e}")
    finally:
        manager.disconnect(websocket)
        if websocket.client_state != WebSocketState.DISCONNECTED:
            await websocket.close()




@router.on_event("startup")
async def startup_event():
    """Запуск таска для трансляции сетевых данных при старте приложения.

    Создаёт асинхронную задачу для периодической трансляции данных.

    Raises:
        Exception: Если произошла ошибка при запуске таска.
    """
    asyncio.create_task(broadcast_data())




async def broadcast_data(translations: Translations = Depends(get_translations())):
    """Таск для периодической трансляции сетевых данных.

    Выполняет сканирование сети и отправляет данные всем подключённым клиентам
    с интервалом в 2 секунды.

    Args:
        translations: Объект переводов для текущего языка.

    Raises:
        Exception: Если произошла ошибка при получении данных или трансляции.
    """
    while True:
        try:
            network_data = await get_network_data()
            if "activation" not in network_data:
                logger.error(translations.gettext("Ошибка: Ключ 'activation' отсутствует в полученных данных сети"))
                await asyncio.sleep(2)
                continue

            await manager.broadcast(network_data)
        except Exception as e:
            logger.error(f"Error in broadcast_data: {e}")
        await asyncio.sleep(2)