malenia/services/rw.py

import logging
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Optional
from uuid import UUID

from remnawave import RemnawaveSDK
from remnawave.models import UpdateUserRequestDto

logger = logging.getLogger(__name__)


# ─────────────────────────────────────────────────────────────────
#  Dataclass — унифицированное представление пользователя RW
# ─────────────────────────────────────────────────────────────────


@dataclass
class RWUserInfo:
    uuid: str
    username: str
    status: str  # "ACTIVE", "DISABLED", "EXPIRED", "ON_HOLD", "LIMITED"
    expire_at: datetime
    used_traffic_bytes: float
    traffic_limit_bytes: int  # 0 = unlimited
    hwid_device_limit: Optional[int]  # None = unlimited
    active_squads: list[dict]  # [{"uuid": "...", "name": "..."}]
    description: Optional[str]
    short_uuid: str


# ─────────────────────────────────────────────────────────────────
#  Внутренние хелперы
# ─────────────────────────────────────────────────────────────────


def _parse_user(user_dto) -> RWUserInfo:
    """Преобразовать UserResponseDto → RWUserInfo."""
    active_squads: list[dict] = []
    for squad in user_dto.active_internal_squads or []:
        active_squads.append({"uuid": str(squad.uuid), "name": squad.name})

    # UserStatus может быть enum-объектом или строкой — нормализуем
    status_raw = user_dto.status
    if hasattr(status_raw, "value"):
        status_str = str(status_raw.value).upper()
    else:
        status_str = str(status_raw).upper()

    # Трафик лежит в user_traffic.used_traffic_bytes
    used_traffic: float = 0.0
    if user_dto.user_traffic is not None:
        used_traffic = float(user_dto.user_traffic.used_traffic_bytes or 0)

    return RWUserInfo(
        uuid=str(user_dto.uuid),
        username=user_dto.username,
        status=status_str,
        expire_at=user_dto.expire_at,
        used_traffic_bytes=used_traffic,
        traffic_limit_bytes=int(user_dto.traffic_limit_bytes or 0),
        hwid_device_limit=user_dto.hwid_device_limit,
        active_squads=active_squads,
        description=user_dto.description,
        short_uuid=user_dto.short_uuid,
    )


# ─────────────────────────────────────────────────────────────────
#  Публичные функции — обёртки с перехватом ошибок
# ─────────────────────────────────────────────────────────────────


async def get_user_by_telegram_id(
    sdk: RemnawaveSDK,
    telegram_id: int,
) -> Optional[RWUserInfo]:
    """
    Получить первого пользователя Remnawave, привязанного к Telegram ID.
    Возвращает None, если пользователь не найден или Remnawave недоступен.
    """
    try:
        users_list = await sdk.users.get_users_by_telegram_id(str(telegram_id))
        if not users_list:
            return None
        return _parse_user(users_list[0])
    except Exception as exc:
        logger.warning(
            "Failed to get RW user by telegram_id=%s: %s",
            telegram_id,
            exc,
        )
        return None


async def get_all_squads(sdk: RemnawaveSDK) -> list[dict]:
    """
    Получить все Internal Squads из Remnawave.
    Возвращает список [{"uuid": "...", "name": "..."}].
    При ошибке возвращает пустой список.
    """
    try:
        response = await sdk.internal_squads.get_internal_squads()
        result: list[dict] = []
        # GetAllInternalSquadsResponseDto имеет поле .internal_squads
        for squad in response.internal_squads:
            result.append({"uuid": str(squad.uuid), "name": squad.name})
        return result
    except Exception as exc:
        logger.warning("Failed to get all internal squads: %s", exc)
        return []


async def add_days(
    sdk: RemnawaveSDK,
    user_uuid: str,
    days: int,
) -> Optional[datetime]:
    """
    Добавить N дней к expire_at пользователя.
    Возвращает новую дату окончания или None при ошибке.
    """
    try:
        # Получаем актуальный expire_at — нельзя слепо использовать
        # закешированное значение, чтобы не потерять изменения других агентов
        user_dto = await sdk.users.get_user_by_uuid(uuid=user_uuid)
        if user_dto is None or user_dto.expire_at is None:
            logger.warning("Cannot add days: user %s has no expire_at", user_uuid)
            return None

        # Обеспечиваем timezone-aware datetime
        current_expire = user_dto.expire_at
        if current_expire.tzinfo is None:
            current_expire = current_expire.replace(tzinfo=timezone.utc)

        new_expire = current_expire + timedelta(days=days)

        result = await sdk.users.update_user(
            UpdateUserRequestDto(
                uuid=UUID(user_uuid),
                expire_at=new_expire,
            )
        )
        if result is None:
            return None

        new_dt = result.expire_at
        if new_dt is not None and new_dt.tzinfo is None:
            new_dt = new_dt.replace(tzinfo=timezone.utc)
        return new_dt
    except Exception as exc:
        logger.warning("Failed to add %d days for user %s: %s", days, user_uuid, exc)
        return None


async def remove_days(
    sdk: RemnawaveSDK,
    user_uuid: str,
    days: int,
) -> Optional[datetime]:
    """
    Вычесть N дней из expire_at пользователя.
    Возвращает новую дату окончания или None при ошибке.
    """
    try:
        user_dto = await sdk.users.get_user_by_uuid(uuid=user_uuid)
        if user_dto is None or user_dto.expire_at is None:
            logger.warning("Cannot remove days: user %s has no expire_at", user_uuid)
            return None

        current_expire = user_dto.expire_at
        if current_expire.tzinfo is None:
            current_expire = current_expire.replace(tzinfo=timezone.utc)

        new_expire = current_expire - timedelta(days=days)

        result = await sdk.users.update_user(
            UpdateUserRequestDto(
                uuid=UUID(user_uuid),
                expire_at=new_expire,
            )
        )
        if result is None:
            return None

        new_dt = result.expire_at
        if new_dt is not None and new_dt.tzinfo is None:
            new_dt = new_dt.replace(tzinfo=timezone.utc)
        return new_dt
    except Exception as exc:
        logger.warning("Failed to remove %d days for user %s: %s", days, user_uuid, exc)
        return None


async def set_hwid_limit(
    sdk: RemnawaveSDK,
    user_uuid: str,
    limit: int,
) -> bool:
    """
    Установить лимит HWID-устройств.
    limit == 0 → убрать ограничение (None).
    Возвращает True при успехе, False при ошибке.
    """
    try:
        hwid_value: Optional[int] = None if limit == 0 else limit
        await sdk.users.update_user(
            UpdateUserRequestDto(
                uuid=UUID(user_uuid),
                hwid_device_limit=hwid_value,
            )
        )
        return True
    except Exception as exc:
        logger.warning(
            "Failed to set HWID limit=%s for user %s: %s",
            limit,
            user_uuid,
            exc,
        )
        return False


async def set_description(
    sdk: RemnawaveSDK,
    user_uuid: str,
    note: str,
) -> bool:
    """
    Установить заметку (description) пользователя в Remnawave.
    Возвращает True при успехе, False при ошибке.
    """
    try:
        await sdk.users.update_user(
            UpdateUserRequestDto(
                uuid=UUID(user_uuid),
                description=note,
            )
        )
        return True
    except Exception as exc:
        logger.warning(
            "Failed to set description for user %s: %s",
            user_uuid,
            exc,
        )
        return False


async def reset_traffic(
    sdk: RemnawaveSDK,
    user_uuid: str,
) -> bool:
    """
    Сбросить счётчик использованного трафика пользователя.
    Возвращает True при успехе, False при ошибке.
    """
    try:
        await sdk.users.reset_user_traffic(uuid=user_uuid)
        return True
    except Exception as exc:
        logger.warning("Failed to reset traffic for user %s: %s", user_uuid, exc)
        return False


async def disable_user(
    sdk: RemnawaveSDK,
    user_uuid: str,
) -> bool:
    """
    Отключить пользователя (статус → DISABLED).
    Возвращает True при успехе, False при ошибке.
    """
    try:
        await sdk.users.disable_user(uuid=user_uuid)
        return True
    except Exception as exc:
        logger.warning("Failed to disable user %s: %s", user_uuid, exc)
        return False


async def enable_user(
    sdk: RemnawaveSDK,
    user_uuid: str,
) -> bool:
    """
    Включить пользователя (статус → ACTIVE).
    Возвращает True при успехе, False при ошибке.
    """
    try:
        await sdk.users.enable_user(uuid=user_uuid)
        return True
    except Exception as exc:
        logger.warning("Failed to enable user %s: %s", user_uuid, exc)
        return False


async def update_squads(
    sdk: RemnawaveSDK,
    user_uuid: str,
    squad_uuid_list: list[str],
) -> bool:
    """
    Полностью заменить список Internal Squads пользователя.
    squad_uuid_list — строковые UUID всех squad'ов (полный новый список).
    Возвращает True при успехе, False при ошибке.
    """
    try:
        await sdk.users.update_user(
            UpdateUserRequestDto(
                uuid=UUID(user_uuid),
                active_internal_squads=[UUID(s) for s in squad_uuid_list],
            )
        )
        return True
    except Exception as exc:
        logger.warning(
            "Failed to update squads for user %s: %s",
            user_uuid,
            exc,
        )
        return False