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

import logging
from typing import Any, Dict, List, Optional, Tuple

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

from app.database.models import PromoGroup, User


def _normalize_period_discounts(period_discounts: Optional[Dict[int, int]]) -> Dict[int, int]:
    if not period_discounts:
        return {}

    normalized: Dict[int, int] = {}

    for key, value in period_discounts.items():
        try:
            period = int(key)
            percent = int(value)
        except (TypeError, ValueError):
            continue

        normalized[period] = max(0, min(100, percent))

    return normalized

logger = logging.getLogger(__name__)

_UNSET = object()


def _normalize_auto_assign_amount(amount_kopeks: Optional[int]) -> Optional[int]:
    if amount_kopeks is None:
        return None

    try:
        normalized = int(amount_kopeks)
    except (TypeError, ValueError):
        return None

    if normalized <= 0:
        return None

    return normalized


async def get_promo_groups_with_counts(
    db: AsyncSession,
) -> List[Tuple[PromoGroup, int]]:
    result = await db.execute(
        select(PromoGroup, func.count(User.id))
        .outerjoin(User, User.promo_group_id == PromoGroup.id)
        .group_by(PromoGroup.id)
        .order_by(PromoGroup.is_default.desc(), PromoGroup.name)
    )
    return result.all()


async def get_promo_group_by_id(db: AsyncSession, group_id: int) -> Optional[PromoGroup]:
    return await db.get(PromoGroup, group_id)


async def get_default_promo_group(db: AsyncSession) -> Optional[PromoGroup]:
    result = await db.execute(
        select(PromoGroup).where(PromoGroup.is_default.is_(True))
    )
    return result.scalars().first()


async def create_promo_group(
    db: AsyncSession,
    name: str,
    *,
    server_discount_percent: int,
    traffic_discount_percent: int,
    device_discount_percent: int,
    period_discounts: Optional[Dict[int, int]] = None,
    auto_assign_amount_kopeks: Optional[int] = None,
) -> PromoGroup:
    normalized_period_discounts = _normalize_period_discounts(period_discounts)
    normalized_auto_amount = _normalize_auto_assign_amount(auto_assign_amount_kopeks)

    promo_group = PromoGroup(
        name=name.strip(),
        server_discount_percent=max(0, min(100, server_discount_percent)),
        traffic_discount_percent=max(0, min(100, traffic_discount_percent)),
        device_discount_percent=max(0, min(100, device_discount_percent)),
        period_discounts=normalized_period_discounts or None,
        auto_assign_amount_kopeks=normalized_auto_amount,
        is_default=False,
    )

    db.add(promo_group)
    await db.commit()
    await db.refresh(promo_group)

    logger.info(
        "Создана промогруппа '%s' с скидками (servers=%s%%, traffic=%s%%, devices=%s%%, periods=%s)",
        promo_group.name,
        promo_group.server_discount_percent,
        promo_group.traffic_discount_percent,
        promo_group.device_discount_percent,
        normalized_period_discounts,
    )

    return promo_group


async def update_promo_group(
    db: AsyncSession,
    group: PromoGroup,
    *,
    name: Optional[str] = None,
    server_discount_percent: Optional[int] = None,
    traffic_discount_percent: Optional[int] = None,
    device_discount_percent: Optional[int] = None,
    period_discounts: Optional[Dict[int, int]] = None,
    auto_assign_amount_kopeks: Any = _UNSET,
) -> PromoGroup:
    if name is not None:
        group.name = name.strip()
    if server_discount_percent is not None:
        group.server_discount_percent = max(0, min(100, server_discount_percent))
    if traffic_discount_percent is not None:
        group.traffic_discount_percent = max(0, min(100, traffic_discount_percent))
    if device_discount_percent is not None:
        group.device_discount_percent = max(0, min(100, device_discount_percent))
    if period_discounts is not None:
        normalized_period_discounts = _normalize_period_discounts(period_discounts)
        group.period_discounts = normalized_period_discounts or None
    if auto_assign_amount_kopeks is not _UNSET:
        group.auto_assign_amount_kopeks = _normalize_auto_assign_amount(auto_assign_amount_kopeks)

    await db.commit()
    await db.refresh(group)

    logger.info(
        "Обновлена промогруппа '%s' (id=%s)",
        group.name,
        group.id,
    )
    return group


async def delete_promo_group(db: AsyncSession, group: PromoGroup) -> bool:
    if group.is_default:
        logger.warning("Попытка удалить базовую промогруппу запрещена")
        return False

    default_group = await get_default_promo_group(db)
    if not default_group:
        logger.error("Не найдена базовая промогруппа для reassignment")
        return False

    await db.execute(
        update(User)
        .where(User.promo_group_id == group.id)
        .values(promo_group_id=default_group.id)
    )
    await db.delete(group)
    await db.commit()

    logger.info(
        "Промогруппа '%s' (id=%s) удалена, пользователи переведены в '%s'",
        group.name,
        group.id,
        default_group.name,
    )
    return True


async def get_promo_group_members(
    db: AsyncSession,
    group_id: int,
    *,
    offset: int = 0,
    limit: int = 20,
) -> List[User]:
    result = await db.execute(
        select(User)
        .options(selectinload(User.subscription))
        .where(User.promo_group_id == group_id)
        .order_by(User.created_at.desc())
        .offset(offset)
        .limit(limit)
    )
    return result.scalars().all()


async def count_promo_group_members(db: AsyncSession, group_id: int) -> int:
    result = await db.execute(
        select(func.count(User.id)).where(User.promo_group_id == group_id)
    )
    return result.scalar_one()


async def get_auto_assign_promo_group(
    db: AsyncSession,
    total_amount_kopeks: int,
) -> Optional[PromoGroup]:
    if total_amount_kopeks <= 0:
        return None

    result = await db.execute(
        select(PromoGroup)
        .where(
            PromoGroup.auto_assign_amount_kopeks.is_not(None),
            PromoGroup.auto_assign_amount_kopeks > 0,
            PromoGroup.auto_assign_amount_kopeks <= total_amount_kopeks,
        )
        .order_by(desc(PromoGroup.auto_assign_amount_kopeks), PromoGroup.id)
    )

    return result.scalars().first()