""" Referral rewards (3C.2). Goals: - Deterministic, abuse-resistant rewards - No manual state tracking per referral; we compute from authoritative DB state - Idempotent updates (can be run via scheduler and on-demand) Current reward: - For every N qualified referrals, grant +M bonus watchlist domain slots. Qualified referral definition: - referred user has `users.referred_by_user_id = referrer.id` - referred user is_active AND is_verified - referred user has an active subscription that is NOT Scout (Trader/Tycoon), and is currently active """ from __future__ import annotations from dataclasses import dataclass from datetime import datetime, timedelta from typing import Optional from sqlalchemy import and_, func, or_, select from sqlalchemy.ext.asyncio import AsyncSession from app.config import get_settings from app.models.subscription import Subscription, SubscriptionStatus, SubscriptionTier from app.models.telemetry import TelemetryEvent from app.models.user import User QUALIFIED_REFERRAL_BATCH_SIZE = 3 BONUS_DOMAINS_PER_BATCH = 5 settings = get_settings() def compute_bonus_domains(qualified_referrals: int) -> int: if qualified_referrals <= 0: return 0 batches = qualified_referrals // QUALIFIED_REFERRAL_BATCH_SIZE return int(batches * BONUS_DOMAINS_PER_BATCH) def compute_badge(qualified_referrals: int) -> Optional[str]: if qualified_referrals >= 10: return "elite_referrer" if qualified_referrals >= 3: return "verified_referrer" return None @dataclass(frozen=True) class ReferralRewardSnapshot: referrer_user_id: int referred_users_total: int qualified_referrals_total: int cooldown_days: int disqualified_cooldown_total: int disqualified_missing_ip_total: int disqualified_shared_ip_total: int disqualified_duplicate_ip_total: int bonus_domains: int badge: Optional[str] computed_at: datetime async def get_referral_reward_snapshot(db: AsyncSession, referrer_user_id: int) -> ReferralRewardSnapshot: # Total referred users (all-time) referred_users_total = int( ( await db.execute( select(func.count(User.id)).where(User.referred_by_user_id == referrer_user_id) ) ).scalar() or 0 ) now = datetime.utcnow() cooldown_days = max(0, int(getattr(settings, "referral_rewards_cooldown_days", 7) or 0)) cooldown_cutoff = now - timedelta(days=cooldown_days) if cooldown_days else None # Referrer IP hashes (window) for self-ref/shared-ip checks ip_window_days = max(1, int(getattr(settings, "referral_rewards_ip_window_days", 30) or 30)) ip_window_start = now - timedelta(days=ip_window_days) referrer_ip_rows = ( await db.execute( select(TelemetryEvent.ip_hash) .where( and_( TelemetryEvent.user_id == referrer_user_id, TelemetryEvent.ip_hash.isnot(None), TelemetryEvent.created_at >= ip_window_start, TelemetryEvent.created_at <= now, ) ) .distinct() ) ).all() referrer_ip_hashes = {str(r[0]) for r in referrer_ip_rows if r and r[0]} # Referred user's registration IP hash (from telemetry) as subquery reg_ip_subq = ( select( TelemetryEvent.user_id.label("user_id"), func.max(TelemetryEvent.ip_hash).label("signup_ip_hash"), ) .where( and_( TelemetryEvent.event_name == "user_registered", TelemetryEvent.user_id.isnot(None), ) ) .group_by(TelemetryEvent.user_id) .subquery() ) # Candidate referred users (paid + verified + active) rows = ( await db.execute( select( User.id, User.created_at, Subscription.started_at, reg_ip_subq.c.signup_ip_hash, ) .select_from(User) .join(Subscription, Subscription.user_id == User.id) .outerjoin(reg_ip_subq, reg_ip_subq.c.user_id == User.id) .where( and_( User.referred_by_user_id == referrer_user_id, User.is_active == True, User.is_verified == True, Subscription.tier.in_([SubscriptionTier.TRADER, SubscriptionTier.TYCOON]), Subscription.status.in_([SubscriptionStatus.ACTIVE, SubscriptionStatus.PAST_DUE]), or_(Subscription.expires_at.is_(None), Subscription.expires_at >= now), ) ) ) ).all() require_ip = bool(getattr(settings, "referral_rewards_require_ip_hash", True)) disqualified_cooldown_total = 0 disqualified_missing_ip_total = 0 disqualified_shared_ip_total = 0 disqualified_duplicate_ip_total = 0 qualified_ip_hashes: set[str] = set() qualified_referrals_total = 0 for _user_id, user_created_at, sub_started_at, signup_ip_hash in rows: # Cooldown: user account age AND subscription age must pass cooldown if cooldown_cutoff is not None: if (user_created_at and user_created_at > cooldown_cutoff) or ( sub_started_at and sub_started_at > cooldown_cutoff ): disqualified_cooldown_total += 1 continue ip_hash = str(signup_ip_hash) if signup_ip_hash else None if require_ip and not ip_hash: disqualified_missing_ip_total += 1 continue if ip_hash and referrer_ip_hashes and ip_hash in referrer_ip_hashes: disqualified_shared_ip_total += 1 continue if ip_hash and ip_hash in qualified_ip_hashes: disqualified_duplicate_ip_total += 1 continue if ip_hash: qualified_ip_hashes.add(ip_hash) qualified_referrals_total += 1 bonus_domains = compute_bonus_domains(qualified_referrals_total) badge = compute_badge(qualified_referrals_total) return ReferralRewardSnapshot( referrer_user_id=referrer_user_id, referred_users_total=referred_users_total, qualified_referrals_total=qualified_referrals_total, cooldown_days=cooldown_days, disqualified_cooldown_total=disqualified_cooldown_total, disqualified_missing_ip_total=disqualified_missing_ip_total, disqualified_shared_ip_total=disqualified_shared_ip_total, disqualified_duplicate_ip_total=disqualified_duplicate_ip_total, bonus_domains=bonus_domains, badge=badge, computed_at=datetime.utcnow(), ) async def apply_referral_rewards_for_user(db: AsyncSession, referrer_user_id: int) -> ReferralRewardSnapshot: """ Apply rewards to the referrer's subscription row, based on current qualified referrals. This is idempotent: it sets the bonus to the computed value. """ snapshot = await get_referral_reward_snapshot(db, referrer_user_id) sub_res = await db.execute(select(Subscription).where(Subscription.user_id == referrer_user_id)) sub = sub_res.scalar_one_or_none() if not sub: # Create default subscription so bonus can be stored sub = Subscription(user_id=referrer_user_id, tier=SubscriptionTier.SCOUT, max_domains=5) db.add(sub) await db.flush() desired = int(snapshot.bonus_domains) current = int(getattr(sub, "referral_bonus_domains", 0) or 0) if current != desired: sub.referral_bonus_domains = desired await db.flush() return snapshot async def apply_referral_rewards_all(db: AsyncSession) -> dict[str, int]: """ Apply rewards for all users that have an invite_code. """ res = await db.execute(select(User.id).where(User.invite_code.isnot(None))) user_ids = [int(r[0]) for r in res.all()] updated = 0 processed = 0 for user_id in user_ids: processed += 1 snap = await get_referral_reward_snapshot(db, user_id) sub_res = await db.execute(select(Subscription).where(Subscription.user_id == user_id)) sub = sub_res.scalar_one_or_none() if not sub: sub = Subscription(user_id=user_id, tier=SubscriptionTier.SCOUT, max_domains=5) db.add(sub) await db.flush() desired = int(snap.bonus_domains) current = int(getattr(sub, "referral_bonus_domains", 0) or 0) if current != desired: sub.referral_bonus_domains = desired updated += 1 return {"processed": processed, "updated": updated}