read.markets/app/services/referral_service.py

"""Referral-code generation, lookup, and signup-time linkage.

D.1 lays down the bookkeeping only — actual credit application happens
in D.3 when the Paddle webhook fires. 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 new user's first paid subscription (D.3), we read the
   `Referral` row to apply discounts to both parties.

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 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


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


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