remnawave-minishop/db/dal/active_discount_dal.py

import logging
from typing import Optional, List
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from sqlalchemy import delete
from datetime import datetime, timezone

from db.models import ActiveDiscount


async def set_active_discount(
    session: AsyncSession,
    user_id: int,
    promo_code_id: int,
    discount_percentage: int,
    expires_at: datetime,
) -> Optional[ActiveDiscount]:
    """
    Set active discount for user.
    Returns None if user already has an active discount (enforce one-at-a-time rule).
    """
    now_utc = datetime.now(timezone.utc)

    existing = await get_active_discount(session, user_id, include_expired=True)
    if existing and existing.expires_at > now_utc:
        logging.warning(
            f"User {user_id} already has active discount (promo_code_id: {existing.promo_code_id}). "
            f"Cannot activate new discount {promo_code_id}."
        )
        return None

    if existing and existing.expires_at <= now_utc:
        await clear_active_discount_if_expired(session, user_id, now=now_utc)

    # Create new active discount
    new_discount = ActiveDiscount(
        user_id=user_id,
        promo_code_id=promo_code_id,
        discount_percentage=discount_percentage,
        activated_at=now_utc,
        expires_at=expires_at,
    )
    session.add(new_discount)
    await session.flush()
    await session.refresh(new_discount)
    logging.info(
        f"Active discount set for user {user_id}: promo_code_id={promo_code_id}, "
        f"discount={discount_percentage}%"
    )
    return new_discount


async def get_active_discount(
    session: AsyncSession,
    user_id: int,
    include_expired: bool = False,
) -> Optional[ActiveDiscount]:
    """Get active discount for user if exists."""
    now_utc = datetime.now(timezone.utc)
    stmt = select(ActiveDiscount).where(ActiveDiscount.user_id == user_id)
    if not include_expired:
        stmt = stmt.where(ActiveDiscount.expires_at > now_utc)
    result = await session.execute(stmt)
    return result.scalar_one_or_none()


async def clear_active_discount(
    session: AsyncSession,
    user_id: int
) -> bool:
    """
    Clear active discount for user.
    Returns True if discount was cleared, False if no discount was found.
    """
    stmt = delete(ActiveDiscount).where(ActiveDiscount.user_id == user_id)
    result = await session.execute(stmt)
    await session.flush()
    cleared = result.rowcount > 0
    if cleared:
        logging.info(f"Active discount cleared for user {user_id}")
    return cleared


async def clear_active_discount_if_expired(
    session: AsyncSession,
    user_id: int,
    now: Optional[datetime] = None,
) -> bool:
    """
    Clear active discount for user only when it has already expired.
    """
    now_utc = now or datetime.now(timezone.utc)
    stmt = delete(ActiveDiscount).where(
        ActiveDiscount.user_id == user_id,
        ActiveDiscount.expires_at <= now_utc,
    )
    result = await session.execute(stmt)
    await session.flush()
    cleared = result.rowcount > 0
    if cleared:
        logging.info("Expired active discount cleared for user %s", user_id)
    return cleared


async def clear_active_discount_if_matches(
    session: AsyncSession,
    user_id: int,
    promo_code_id: Optional[int] = None,
    expires_at_lte: Optional[datetime] = None,
) -> bool:
    """
    Clear active discount for user only when additional constraints match.
    """
    conditions = [ActiveDiscount.user_id == user_id]
    if promo_code_id is not None:
        conditions.append(ActiveDiscount.promo_code_id == promo_code_id)
    if expires_at_lte is not None:
        conditions.append(ActiveDiscount.expires_at <= expires_at_lte)

    stmt = delete(ActiveDiscount).where(*conditions)
    result = await session.execute(stmt)
    await session.flush()
    cleared = result.rowcount > 0
    if cleared:
        logging.info(
            "Active discount cleared for user %s by constrained cleanup.",
            user_id,
        )
    return cleared


async def get_expired_active_discounts(
    session: AsyncSession,
    now: Optional[datetime] = None,
    limit: int = 100,
) -> List[ActiveDiscount]:
    """Get expired active discount reservations for cleanup/notifications."""
    now_utc = now or datetime.now(timezone.utc)
    stmt = (
        select(ActiveDiscount)
        .where(ActiveDiscount.expires_at <= now_utc)
        .order_by(ActiveDiscount.expires_at.asc())
        .limit(limit)
    )
    result = await session.execute(stmt)
    return list(result.scalars().all())


async def clear_active_discounts_by_promo_code(
    session: AsyncSession,
    promo_code_id: int
) -> int:
    """
    Clear all active discounts associated with a specific promo code.
    Returns the number of discounts cleared.
    """
    stmt = delete(ActiveDiscount).where(ActiveDiscount.promo_code_id == promo_code_id)
    result = await session.execute(stmt)
    await session.flush()
    count = result.rowcount
    if count > 0:
        logging.info(f"Cleared {count} active discount(s) for promo_code_id={promo_code_id}")
    return count