"""Referral-code generation, lookup, signup-time linkage, and conversion-time credit grants. The flow: 1. /login renders an "invited" banner when the URL carries `?ref=`. 2. The code travels through the email-OTP flow inside the pending cookie so it survives the GET /login → POST /login → /verify hops. 3. When the new user's row is first created (POST /login on an unknown email), `referred_by_user_id` is set and a `Referral` row is written. 4. On the referred user's first paid subscription, `convert_referral` is called from the Stripe webhook: both parties get a credit-window extension worth the promised "50% off for 3 months" (= 45 days of full paid access via `users.credit_until`), and the Referral row's `converted_at` + `credited_at` are stamped. The code itself is 8 characters from an unambiguous alphabet so users can read it off a phone screen or dictate it over the phone. """ from __future__ import annotations import secrets from datetime import timedelta from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.db import utcnow from app.logging import get_logger from app.models import Referral, User from app.services.access import _aware log = get_logger("referral") # Unambiguous alphabet — no 0/O, no 1/I/L. 32 chars → 8 positions ≈ 1e12 # combinations, plenty for our scale, and a unique-constraint catches # collisions if we ever generate the same one twice. _ALPHABET = "ABCDEFGHJKMNPQRSTUVWXYZ23456789" _CODE_LEN = 8 # Value-equivalent of the public-facing "50% off for 3 months" promise, # delivered as a credit-window extension. 50% × 3 months ≈ 1.5 months # of free service ≈ 45 days. Pure-credit delivery means the mechanism # is processor-agnostic and stacks cleanly when both parties refer. REFERRAL_CREDIT_DAYS = 45 def generate_code() -> str: """Cryptographically random 8-char code from the unambiguous alphabet.""" return "".join(secrets.choice(_ALPHABET) for _ in range(_CODE_LEN)) def normalise_code(raw: str | None) -> str | None: """Trim, uppercase, strip non-alphabet characters. Used on inbound `?ref=` params so users can paste with spaces / lowercase. Returns None if the result isn't a plausible code.""" if not raw: return None cleaned = "".join(c for c in raw.upper() if c in _ALPHABET) if len(cleaned) != _CODE_LEN: return None return cleaned async def assign_code_if_missing(session: AsyncSession, user: User) -> User: """Generate + persist a referral code on `user` if they don't have one yet. Retries on the (very rare) collision. The `user` argument is the User attached to the auth-dependency session, which has since been closed — so it is detached from our `session`. We re-fetch it here before mutating so SQLAlchemy doesn't refuse with 'not persistent within this Session'. """ if user.referral_code: return user db_user = await session.get(User, user.id) if db_user is None: raise RuntimeError(f"referral_service: user {user.id} vanished mid-request") if db_user.referral_code: # Raced with another request — accept their code. return db_user for _ in range(8): code = generate_code() existing = (await session.execute( select(User.id).where(User.referral_code == code) )).scalar_one_or_none() if existing is None: db_user.referral_code = code await session.commit() log.info("referral.code_assigned", user_id=db_user.id, code=code) return db_user # 8 collisions in a row would be a statistical event we'd want to # know about. raise RuntimeError("referral_service: exhausted code-collision retries") async def lookup_referrer(session: AsyncSession, code: str | None) -> User | None: """Return the User whose `referral_code` matches, or None. Normalises the input via `normalise_code` so URL-paste variations all resolve.""" code = normalise_code(code) if not code: return None return (await session.execute( select(User).where(User.referral_code == code) )).scalar_one_or_none() async def link_new_user( session: AsyncSession, new_user: User, referrer: User | None, ) -> Referral | None: """Record a referral if the supplied referrer is valid. Idempotent (safe to call multiple times for the same new user — the unique constraint on `referred_user_id` makes duplicate inserts a no-op). Self-referral is silently rejected. """ if referrer is None or new_user.id is None or referrer.id == new_user.id: return None if new_user.referred_by_user_id is not None: # Already linked; this user can't be referred twice. return None new_user.referred_by_user_id = referrer.id ref = Referral( referrer_user_id=referrer.id, referred_user_id=new_user.id, created_at=utcnow(), ) session.add(ref) await session.commit() await session.refresh(new_user) await session.refresh(ref) log.info( "referral.linked", referrer_id=referrer.id, referred_id=new_user.id, ) return ref def _extend_credit(user: User, days: int) -> None: """Stack `days` of paid-tier credit onto `user.credit_until`. Anchors at max(now, current credit_until) so granting twice gives twice the runway — never shortens the window. Mirrors the cli.grant_credit anchoring rule so manual + automatic grants compose.""" now = utcnow() anchor = max(now, _aware(user.credit_until) or now) user.credit_until = anchor + timedelta(days=days) async def convert_referral( session: AsyncSession, referred_user: User, ) -> Referral | None: """Stamp the Referral row for `referred_user` as converted and grant both parties their credit. Idempotent — safe to call from every subscription event: - Returns None if no Referral row exists for this user (direct signup, no inviter). - Returns the existing Referral (unchanged) if `converted_at` is already set — this is a renewal or duplicate webhook delivery. - Otherwise: extends both users' `credit_until` by REFERRAL_CREDIT_DAYS and sets `converted_at` + `credited_at`. The caller is responsible for committing the session — this lets the Stripe webhook compose the conversion inside its outer audit-row transaction, so a mid-flight failure rolls back the tier flip AND the conversion together. Self-referral cannot happen here in practice (link_new_user blocks it at signup) but we guard anyway: if the row somehow names the same user on both sides, we stamp the timestamps but only credit once.""" row = (await session.execute( select(Referral).where(Referral.referred_user_id == referred_user.id) )).scalar_one_or_none() if row is None: return None if row.converted_at is not None: return row referrer = await session.get(User, row.referrer_user_id) now = utcnow() # Always credit the buyer; credit the referrer too unless they're # the same row (defence-in-depth) or have been deleted. _extend_credit(referred_user, REFERRAL_CREDIT_DAYS) if referrer is not None and referrer.id != referred_user.id: _extend_credit(referrer, REFERRAL_CREDIT_DAYS) row.converted_at = now row.credited_at = now log.info( "referral.converted", referral_id=row.id, referrer_id=row.referrer_user_id, referred_id=row.referred_user_id, credit_days=REFERRAL_CREDIT_DAYS, ) return row