Solo_bot/database/referrals.py

from sqlalchemy import and_, desc, func, insert, select, text, update
from sqlalchemy.ext.asyncio import AsyncSession

from config import CHECK_REFERRAL_REWARD_ISSUED, REFERRAL_BONUS_PERCENTAGES
from core.bootstrap import BUTTONS_CONFIG
from database.access.resolution import resolve_user_optional
from database.models import Referral
from logger import logger


async def add_referral(session: AsyncSession, referred_legacy: int, referrer_legacy: int):
    ru = await resolve_user_optional(session, referred_legacy)
    rf = await resolve_user_optional(session, referrer_legacy)
    if ru is None or rf is None:
        return
    if ru.id == rf.id:
        logger.warning(f"⚠️ Попытка самореферала: {referred_legacy}")
        return

    stmt = insert(Referral).values(
        referred_user_id=ru.id,
        referrer_user_id=rf.id,
        referred_tg_id=ru.tg_id,
        referrer_tg_id=rf.tg_id,
    )
    await session.execute(stmt)
    logger.info(f"✅ Добавлена реферальная связь: {ru.id} → {rf.id}")


async def get_referral_by_referred_id(session: AsyncSession, referred_legacy: int) -> dict | None:
    ru = await resolve_user_optional(session, referred_legacy)
    if ru is None:
        return None
    stmt = select(Referral).where(Referral.referred_user_id == ru.id)
    result = await session.execute(stmt)
    row = result.scalar_one_or_none()
    return dict(row.__dict__) if row else None


async def get_total_referrals(session: AsyncSession, referrer_legacy: int) -> int:
    ru = await resolve_user_optional(session, referrer_legacy)
    if ru is None:
        return 0
    stmt = select(func.count()).select_from(Referral).where(Referral.referrer_user_id == ru.id)
    result = await session.execute(stmt)
    return result.scalar()


async def get_active_referrals(session: AsyncSession, referrer_legacy: int) -> int:
    ru = await resolve_user_optional(session, referrer_legacy)
    if ru is None:
        return 0
    stmt = (
        select(func.count())
        .select_from(Referral)
        .where(
            and_(
                Referral.referrer_user_id == ru.id,
                Referral.reward_issued.is_(True),
            )
        )
    )
    result = await session.execute(stmt)
    return result.scalar()


async def mark_referral_reward_issued(session: AsyncSession, referred_legacy: int):
    ru = await resolve_user_optional(session, referred_legacy)
    if ru is None:
        return
    await session.execute(update(Referral).where(Referral.referred_user_id == ru.id).values(reward_issued=True))


async def get_total_referral_bonus(session: AsyncSession, referrer_legacy: int, max_levels: int) -> float:
    referral_enabled = bool(BUTTONS_CONFIG.get("REFERRAL_BUTTON_ENABLED", True))
    if not referral_enabled:
        logger.debug("Реферальная программа отключена, бонусы не начисляются")
        return 0.0

    ru = await resolve_user_optional(session, referrer_legacy)
    if ru is None:
        return 0.0
    uid = ru.id

    if CHECK_REFERRAL_REWARD_ISSUED:
        bonus_cte = """
            WITH RECURSIVE
            referral_levels AS (
                SELECT 
                    referred_user_id, 
                    referrer_user_id, 
                    1 AS level
                FROM referrals 
                WHERE referrer_user_id = :user_id AND reward_issued = TRUE
                
                UNION
                
                SELECT 
                    r.referred_user_id, 
                    r.referrer_user_id, 
                    rl.level + 1
                FROM referrals r
                JOIN referral_levels rl ON r.referrer_user_id = rl.referred_user_id
                WHERE rl.level < :max_levels AND r.reward_issued = TRUE
            ),
            earliest_payments AS (
                SELECT DISTINCT ON (user_id) user_id, amount, created_at
                FROM payments
                WHERE status = 'success' 
                  AND payment_system NOT IN ('coupon', 'admin', 'referral')
                ORDER BY user_id, created_at
            )
        """
        bonus_query = (
            bonus_cte
            + f"""
            SELECT 
                COALESCE(SUM(
                    CASE
                        {
                " ".join([
                    f"WHEN rl.level = {level} THEN {REFERRAL_BONUS_PERCENTAGES[level]} * ep.amount"
                    if isinstance(REFERRAL_BONUS_PERCENTAGES[level], float)
                    else f"WHEN rl.level = {level} THEN {REFERRAL_BONUS_PERCENTAGES[level]}"
                    for level in REFERRAL_BONUS_PERCENTAGES
                ])
            }
                        ELSE 0 
                    END
                ), 0) AS total_bonus
            FROM referral_levels rl
            JOIN earliest_payments ep ON rl.referred_user_id = ep.user_id
            WHERE rl.level <= :max_levels
        """
        )
    else:
        bonus_cte = """
            WITH RECURSIVE
            referral_levels AS (
                SELECT 
                    referred_user_id, 
                    referrer_user_id, 
                    1 AS level
                FROM referrals 
                WHERE referrer_user_id = :user_id
                
                UNION
                
                SELECT 
                    r.referred_user_id, 
                    r.referrer_user_id, 
                    rl.level + 1
                FROM referrals r
                JOIN referral_levels rl ON r.referrer_user_id = rl.referred_user_id
                WHERE rl.level < :max_levels
            )
        """
        bonus_query = (
            bonus_cte
            + f"""
            SELECT 
                COALESCE(SUM(
                    CASE
                        {
                " ".join([
                    f"WHEN rl.level = {level} THEN {REFERRAL_BONUS_PERCENTAGES[level]} * p.amount"
                    if isinstance(REFERRAL_BONUS_PERCENTAGES[level], float)
                    else f"WHEN rl.level = {level} THEN {REFERRAL_BONUS_PERCENTAGES[level]}"
                    for level in REFERRAL_BONUS_PERCENTAGES
                ])
            }
                        ELSE 0 
                    END
                ), 0) AS total_bonus
            FROM referral_levels rl
            JOIN payments p ON rl.referred_user_id = p.user_id
            WHERE p.status = 'success' 
              AND p.payment_system NOT IN ('coupon', 'admin', 'referral')
              AND rl.level <= :max_levels
        """
        )

    result = await session.execute(
        text(bonus_query),
        {"user_id": uid, "max_levels": max_levels},
    )
    total_bonus_raw = result.scalar()
    total_bonus = round(float(total_bonus_raw or 0), 2)

    logger.debug(f"Получена общая сумма бонусов от рефералов: {total_bonus}")
    return total_bonus


async def get_referrals_by_level(session: AsyncSession, referrer_legacy: int, max_levels: int) -> dict:
    ru = await resolve_user_optional(session, referrer_legacy)
    if ru is None:
        return {}
    query = """
        WITH RECURSIVE referral_levels AS (
            SELECT referred_user_id, referrer_user_id, 1 AS level 
            FROM referrals 
            WHERE referrer_user_id = :referrer_user_id
            UNION
            SELECT r.referred_user_id, r.referrer_user_id, rl.level + 1
            FROM referrals r
            JOIN referral_levels rl ON r.referrer_user_id = rl.referred_user_id
            WHERE rl.level < :max_levels
        )
        SELECT level, 
               COUNT(*) AS level_count, 
               COUNT(CASE WHEN reward_issued THEN 1 END) AS active_level_count
        FROM referral_levels rl
        JOIN referrals r ON rl.referred_user_id = r.referred_user_id
        GROUP BY level
        ORDER BY level
    """
    result = await session.execute(
        text(query),
        {"referrer_user_id": ru.id, "max_levels": max_levels},
    )
    return {
        row["level"]: {
            "total": row["level_count"],
            "active": row["active_level_count"],
        }
        for row in result.mappings()
    }


async def get_referral_stats(session: AsyncSession, referrer_legacy: int):
    logger.info(f"[ReferralStats] Получение статистики для пользователя {referrer_legacy}")

    total_referrals = await get_total_referrals(session, referrer_legacy)
    active_referrals = await get_active_referrals(session, referrer_legacy)
    max_levels = len(REFERRAL_BONUS_PERCENTAGES)
    referrals_by_level = await get_referrals_by_level(session, referrer_legacy, max_levels)
    total_referral_bonus = await get_total_referral_bonus(session, referrer_legacy, max_levels)

    return {
        "total_referrals": total_referrals,
        "active_referrals": active_referrals,
        "referrals_by_level": referrals_by_level,
        "total_referral_bonus": total_referral_bonus,
    }


async def get_user_referral_count(session: AsyncSession, legacy: int) -> int:
    ru = await resolve_user_optional(session, legacy)
    if ru is None:
        return 0
    result = await session.execute(select(func.count()).select_from(Referral).where(Referral.referrer_user_id == ru.id))
    return result.scalar_one() or 0


async def get_referral_position(session: AsyncSession, referral_count: int) -> int:
    subq = (
        select(Referral.referrer_user_id)
        .group_by(Referral.referrer_user_id)
        .having(func.count() > referral_count)
        .subquery()
    )
    query = select(func.count()).select_from(subq)
    result = await session.execute(query)
    count = result.scalar() or 0
    return count + 1


async def get_top_referrals(session: AsyncSession, limit: int = 5):
    query = (
        select(Referral.referrer_user_id, func.count().label("referral_count"))
        .group_by(Referral.referrer_user_id)
        .order_by(desc("referral_count"))
        .limit(limit)
    )
    result = await session.execute(query)
    return [{"referrer_user_id": row.referrer_user_id, "referral_count": row.referral_count} for row in result.all()]