backend/app/services/mongo_rep.py

import os
from motor.motor_asyncio import AsyncIOMotorClient
from pymongo.errors import DuplicateKeyError, NetworkTimeout
import logging


class MongoDBRepository:
    def __init__(self):
        # Настройки MongoDB из переменных окружения
        mongo_uri = os.getenv("MONGO_URL")
        database_name = os.getenv("DB_NAME")
        server_collection = os.getenv("SERVER_COLLECTION", "servers")
        plan_collection = os.getenv("PLAN_COLLECTION", "plans")

        # Подключение к базе данных и коллекциям
        self.client = AsyncIOMotorClient(mongo_uri)
        self.db = self.client[database_name]
        self.collection = self.db[server_collection]  # Коллекция серверов
        self.plans_collection = self.db[plan_collection]  # Коллекция планов
        self.logger = logging.getLogger(__name__)

    async def add_subscription_plan(self, plan_data):
        """Добавляет новый тарифный план в коллекцию."""
        try:
            result = await self.plans_collection.insert_one(plan_data)
            self.logger.debug(f"Тарифный план добавлен с ID: {result.inserted_id}")
            return result.inserted_id
        except DuplicateKeyError:
            self.logger.error("Дублирующий ключ.")
        except NetworkTimeout:
            self.logger.error("Сетевой таймаут.")

    async def get_subscription_plan(self, plan_name):
        """Получает тарифный план по его имени."""
        try:
            plan = await self.plans_collection.find_one({"name": plan_name})
            if plan:
                self.logger.debug(f"Найден тарифный план: {plan}")
            else:
                self.logger.error(f"Тарифный план {plan_name} не найден.")
            return plan
        except DuplicateKeyError:
            self.logger.error("Дублирующий ключ.")
        except NetworkTimeout:
            self.logger.error("Сетевой таймаут.")

    async def add_server(self, server_data):
        """Добавляет новый VPN сервер в коллекцию."""
        try:
            result = await self.collection.insert_one(server_data)
            self.logger.debug(f"VPN сервер добавлен с ID: {result.inserted_id}")
            return result.inserted_id
        except DuplicateKeyError:
            self.logger.error("Дублирующий ключ.")
        except NetworkTimeout:
            self.logger.error("Сетевой таймаут.")

    async def get_server(self, server_name: str):
        """Получает сервер VPN по его ID."""
        try:
            server = await self.collection.find_one({"server.name": server_name})
            if server:
                self.logger.debug(f"Найден VPN сервер: {server}")
            else:
                self.logger.debug(f"VPN сервер с ID {server_name} не найден.")
            return server
        except DuplicateKeyError:
            self.logger.error("Дублирующий ключ.")
        except NetworkTimeout:
            self.logger.error("Сетевой таймаут.")
            
    async def get_server_with_least_clients(self):
        """Возвращает сервер с наименьшим количеством подключенных клиентов."""
        try:
            pipeline = [
                {
                    "$addFields": {
                        "current_clients": {"$size": {"$ifNull": ["$clients", []]}}
                    }
                },
                {
                    "$sort": {"current_clients": 1}
                },
                {
                    "$limit": 1
                }
            ]

            result = await self.collection.aggregate(pipeline).to_list(length=1)
            if result:
                server = result[0]
                self.logger.debug(f"Найден сервер с наименьшим количеством клиентов: {server}")
                return server
            else:
                self.logger.debug("Не найдено серверов.")
                return None
        except DuplicateKeyError:
            self.logger.error("Дублирующий ключ.")
        except NetworkTimeout:
            self.logger.error("Сетевой таймаут.")

    async def update_server(self, server_name, update_data):
        """Обновляет данные VPN сервера."""
        try:
            result = await self.collection.update_one({"server_name": server_name}, {"$set": update_data})
            if result.matched_count > 0:
                self.logger.debug(f"VPN сервер с ID {server_name} обновлен.")
            else:
                self.logger.debug(f"VPN сервер с ID {server_name} не найден.")
            return result.matched_count > 0
        except DuplicateKeyError:
            self.logger.error("Дублирующий ключ.")
        except NetworkTimeout:
            self.logger.error("Сетевой таймаут.")

    async def delete_server(self, server_name):
        """Удаляет VPN сервер по его ID."""
        try:
            result = await self.collection.delete_one({"name": server_name})
            if result.deleted_count > 0:
                self.logger.debug(f"VPN сервер с ID {server_name} удален.")
            else:
                self.logger.debug(f"VPN сервер с ID {server_name} не найден.")
            return result.deleted_count > 0
        except DuplicateKeyError:
            self.logger.error("Дублирующий ключ.")
        except NetworkTimeout:
            self.logger.error("Сетевой таймаут.")

    async def list_servers(self):
        """Возвращает список всех VPN серверов."""
        try:
            servers = await self.collection.find().to_list(length=1000)  # Получить до 1000 серверов (можно настроить)
            self.logger.debug(f"Найдено {len(servers)} VPN серверов.")
            return servers
        except DuplicateKeyError:
            self.logger.error("Дублирующий ключ.")
        except NetworkTimeout:
            self.logger.error("Сетевой таймаут.")


    async def __aenter__(self):
        """
        Метод вызывается при входе в блок with.
        """
        self.logger.debug("Контекстный менеджер: подключение открыто.")
        return self  

    async def __aexit__(self, exc_type, exc_value, traceback):
        """
        Метод вызывается при выходе из блока with.
        """
        await self.close_connection() 
        if exc_type:
            self.logger.error(f"Контекстный менеджер завершён с ошибкой: {exc_value}")
        else:
            self.logger.debug("Контекстный менеджер: подключение закрыто.")