remnawave-bedolaga-telegram-bot/app/database/crud/user_promo_group.py

"""CRUD операции для связи пользователей с промогруппами (Many-to-Many)."""

import logging
from datetime import datetime

from sqlalchemy import and_, desc, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from app.database.models import PromoGroup, User, UserPromoGroup


logger = logging.getLogger(__name__)


async def _sync_user_primary_promo_group(
    db: AsyncSession,
    user_id: int,
) -> None:
    """Синхронизирует колонку users.promo_group_id с primary промогруппой."""

    try:
        result = await db.execute(
            select(UserPromoGroup.promo_group_id)
            .join(PromoGroup, UserPromoGroup.promo_group_id == PromoGroup.id)
            .where(UserPromoGroup.user_id == user_id)
            .order_by(desc(PromoGroup.priority), PromoGroup.id)
        )

        first = result.first()
        new_primary_id = first[0] if first else None

        user = await db.get(User, user_id)
        if not user:
            return

        if user.promo_group_id != new_primary_id:
            user.promo_group_id = new_primary_id
            user.updated_at = datetime.utcnow()

    except Exception as error:
        logger.error(
            'Ошибка синхронизации primary промогруппы пользователя %s: %s',
            user_id,
            error,
        )


async def sync_user_primary_promo_group(
    db: AsyncSession,
    user_id: int,
) -> None:
    """Публичная обертка для синхронизации primary промогруппы пользователя."""

    await _sync_user_primary_promo_group(db, user_id)


async def add_user_to_promo_group(
    db: AsyncSession, user_id: int, promo_group_id: int, assigned_by: str = 'admin'
) -> UserPromoGroup | None:
    """
    Добавляет пользователю промогруппу.

    Args:
        db: Сессия БД
        user_id: ID пользователя
        promo_group_id: ID промогруппы
        assigned_by: Кто назначил ('admin', 'system', 'auto', 'promocode')

    Returns:
        UserPromoGroup или None если уже существует
    """
    try:
        # Проверяем существование связи
        existing = await has_user_promo_group(db, user_id, promo_group_id)
        if existing:
            logger.info(f'Пользователь {user_id} уже имеет промогруппу {promo_group_id}')
            return None

        # Создаем новую связь
        user_promo_group = UserPromoGroup(
            user_id=user_id,
            promo_group_id=promo_group_id,
            assigned_by=assigned_by,
        )
        db.add(user_promo_group)
        await db.flush()

        await _sync_user_primary_promo_group(db, user_id)

        await db.commit()
        await db.refresh(user_promo_group)

        logger.info(f'Пользователю {user_id} добавлена промогруппа {promo_group_id} ({assigned_by})')
        return user_promo_group

    except Exception as error:
        logger.error(f'Ошибка добавления промогруппы пользователю: {error}')
        await db.rollback()
        return None


async def remove_user_from_promo_group(db: AsyncSession, user_id: int, promo_group_id: int) -> bool:
    """
    Удаляет промогруппу у пользователя.

    Args:
        db: Сессия БД
        user_id: ID пользователя
        promo_group_id: ID промогруппы

    Returns:
        True если удалено, False если связи не было
    """
    try:
        result = await db.execute(
            select(UserPromoGroup).where(
                and_(UserPromoGroup.user_id == user_id, UserPromoGroup.promo_group_id == promo_group_id)
            )
        )
        user_promo_group = result.scalar_one_or_none()

        if not user_promo_group:
            logger.warning(f'Связь пользователя {user_id} с промогруппой {promo_group_id} не найдена')
            return False

        await db.delete(user_promo_group)
        await db.flush()

        await _sync_user_primary_promo_group(db, user_id)

        await db.commit()

        logger.info(f'У пользователя {user_id} удалена промогруппа {promo_group_id}')
        return True

    except Exception as error:
        logger.error(f'Ошибка удаления промогруппы у пользователя: {error}')
        await db.rollback()
        return False


async def get_user_promo_groups(db: AsyncSession, user_id: int) -> list[UserPromoGroup]:
    """
    Получает все промогруппы пользователя, отсортированные по приоритету.

    Args:
        db: Сессия БД
        user_id: ID пользователя

    Returns:
        Список UserPromoGroup с загруженными PromoGroup, отсортированный по приоритету DESC
    """
    try:
        result = await db.execute(
            select(UserPromoGroup)
            .options(selectinload(UserPromoGroup.promo_group))
            .where(UserPromoGroup.user_id == user_id)
            .join(PromoGroup, UserPromoGroup.promo_group_id == PromoGroup.id)
            .order_by(desc(PromoGroup.priority), PromoGroup.id)
        )
        return list(result.scalars().all())

    except Exception as error:
        logger.error(f'Ошибка получения промогрупп пользователя {user_id}: {error}')
        return []


async def get_primary_user_promo_group(db: AsyncSession, user_id: int) -> PromoGroup | None:
    """
    Получает промогруппу пользователя с максимальным приоритетом.

    Args:
        db: Сессия БД
        user_id: ID пользователя

    Returns:
        PromoGroup с максимальным приоритетом или None
    """
    try:
        user_promo_groups = await get_user_promo_groups(db, user_id)

        if not user_promo_groups:
            return None

        # Первая в списке имеет максимальный приоритет (список уже отсортирован)
        return user_promo_groups[0].promo_group if user_promo_groups[0].promo_group else None

    except Exception as error:
        logger.error(f'Ошибка получения primary промогруппы пользователя {user_id}: {error}')
        return None


async def has_user_promo_group(db: AsyncSession, user_id: int, promo_group_id: int) -> bool:
    """
    Проверяет наличие промогруппы у пользователя.

    Args:
        db: Сессия БД
        user_id: ID пользователя
        promo_group_id: ID промогруппы

    Returns:
        True если пользователь уже имеет эту промогруппу
    """
    try:
        result = await db.execute(
            select(UserPromoGroup).where(
                and_(UserPromoGroup.user_id == user_id, UserPromoGroup.promo_group_id == promo_group_id)
            )
        )
        return result.scalar_one_or_none() is not None

    except Exception as error:
        logger.error(f'Ошибка проверки промогруппы пользователя: {error}')
        return False


async def count_user_promo_groups(db: AsyncSession, user_id: int) -> int:
    """
    Подсчитывает количество промогрупп у пользователя.

    Args:
        db: Сессия БД
        user_id: ID пользователя

    Returns:
        Количество промогрупп
    """
    try:
        result = await db.execute(select(UserPromoGroup).where(UserPromoGroup.user_id == user_id))
        return len(list(result.scalars().all()))

    except Exception as error:
        logger.error(f'Ошибка подсчета промогрупп пользователя: {error}')
        return 0


async def replace_user_promo_groups(
    db: AsyncSession, user_id: int, promo_group_ids: list[int], assigned_by: str = 'admin'
) -> bool:
    """
    Заменяет все промогруппы пользователя на новый список.

    Args:
        db: Сессия БД
        user_id: ID пользователя
        promo_group_ids: Список ID промогрупп
        assigned_by: Кто назначил

    Returns:
        True если успешно
    """
    try:
        # Удаляем все текущие промогруппы
        await db.execute(select(UserPromoGroup).where(UserPromoGroup.user_id == user_id))
        result = await db.execute(select(UserPromoGroup).where(UserPromoGroup.user_id == user_id))
        for upg in result.scalars().all():
            await db.delete(upg)

        # Добавляем новые
        for promo_group_id in promo_group_ids:
            user_promo_group = UserPromoGroup(user_id=user_id, promo_group_id=promo_group_id, assigned_by=assigned_by)
            db.add(user_promo_group)

        await db.commit()
        logger.info(f'Промогруппы пользователя {user_id} заменены на {promo_group_ids}')
        return True

    except Exception as error:
        logger.error(f'Ошибка замены промогрупп пользователя: {error}')
        await db.rollback()
        return False