read.markets/app/services/referral_service.py

"""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=<code>`.
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=<code>` 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