shaguabijia-app-server/app/repositories/invite.py

"""好友邀请 CRUD(注册即生效,邀请人 + 被邀请人各发金币)。

防重复发奖三道(仿 ad_reward / 提现的资金安全思路):
  1. invitee_user_id 唯一 → 一个被邀请人只能被绑定一次(幂等键)。
  2. 自邀屏蔽 → inviter == invitee 直接拒。
  3. 现成的手机号唯一(每个被邀请人 = 一个真实手机号账号)= 天然限制刷量规模。

发金币复用 wallet.grant_coins(grant 只 flush 不 commit),与建关系记录在**同一事务**
commit,保证"建关系 + 双方加金币"原子。奖励额 = rewards.INVITE_INVITER_COINS /
INVITE_INVITEE_COINS。
"""
from __future__ import annotations

import secrets
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone

from sqlalchemy import func, select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session

from app.core import rewards
from app.models.invite import InviteRelation
from app.models.user import User
from app.repositories import wallet as crud_wallet

# 邀请码字符集:去掉易混字符(0/O/1/I/L/B/8/S/5/Z/2),用户口述/手输不易错
_CODE_ALPHABET = "ACDEFGHJKMNPQRTUVWXY34679"
_CODE_LEN = 6


def _gen_code() -> str:
    return "".join(secrets.choice(_CODE_ALPHABET) for _ in range(_CODE_LEN))


def ensure_code(db: Session, user: User) -> str:
    """保证 user 有邀请码(懒生成),返回它。唯一约束碰撞则换码重试。

    ⚠️ 本函数会 db.commit() 整个 session(与 grant_coins"只 flush 不 commit"约定相反)。
    当前只在 GET /invite/me 里调(该请求此前无其它写,提交范围干净)。若将来在带其它未提交
    写的请求里复用本函数,会被它提前 commit——届时应改成 flush + 由调用方 commit。
    """
    if user.invite_code:
        return user.invite_code
    for _ in range(8):
        user.invite_code = _gen_code()
        try:
            db.commit()
            db.refresh(user)
            return user.invite_code
        except IntegrityError:
            db.rollback()
            user = db.get(User, user.id)  # 重取(rollback 后实例已过期),继续换码
    raise RuntimeError("生成邀请码连续碰撞,请重试")


def resolve_inviter(db: Session, invite_code: str) -> User | None:
    """邀请码 → 邀请人(大小写不敏感)。"""
    code = (invite_code or "").strip().upper()
    if not code:
        return None
    return db.execute(
        select(User).where(User.invite_code == code)
    ).scalar_one_or_none()


def _relation_of_invitee(db: Session, invitee_id: int) -> InviteRelation | None:
    return db.execute(
        select(InviteRelation).where(InviteRelation.invitee_user_id == invitee_id)
    ).scalar_one_or_none()


def _is_new_user(user: User) -> bool:
    """被邀请人是否为"新注册"(created_at 在 INVITE_NEW_USER_WINDOW_HOURS 窗口内)。

    "新用户闸":只奖刚注册的人,挡存量老用户互相填码薅羊毛。兼容 PG(tz-aware)与
    SQLite(naive,按 UTC 解释)。
    """
    created = user.created_at
    if created is None:
        return False
    if created.tzinfo is None:
        created = created.replace(tzinfo=timezone.utc)
    age = datetime.now(timezone.utc) - created
    return age <= timedelta(hours=rewards.INVITE_NEW_USER_WINDOW_HOURS)


@dataclass
class BindResult:
    status: str                          # success / already_bound / invalid_code / self_invite / not_eligible
    relation: InviteRelation | None = None
    invitee_coin: int = 0                # 本次给被邀请人发的金币(success 时 >0)


def bind(
    db: Session, *, invitee: User, invite_code: str, channel: str = "clipboard"
) -> BindResult:
    """把 invitee 绑定到 invite_code 对应的邀请人,注册即生效 + 双方发金币。

    幂等:invitee 已被绑过 → already_bound(不重复发奖)。
    """
    # 幂等:已绑过直接返回(不重复发奖)
    existing = _relation_of_invitee(db, invitee.id)
    if existing is not None:
        return BindResult("already_bound", existing)

    inviter = resolve_inviter(db, invite_code)
    if inviter is None or inviter.status != "active":
        return BindResult("invalid_code")
    if inviter.id == invitee.id:
        return BindResult("self_invite")
    # 新用户闸:被邀请人必须是"新注册"(窗口内)才发奖,挡存量老用户互相填码薅羊毛
    if not _is_new_user(invitee):
        return BindResult("not_eligible")

    inviter_coin = rewards.INVITE_INVITER_COINS
    invitee_coin = rewards.INVITE_INVITEE_COINS

    rel = InviteRelation(
        inviter_user_id=inviter.id,
        invitee_user_id=invitee.id,
        channel=(channel or "clipboard")[:16],
        status="effective",
        inviter_coin=inviter_coin,
        invitee_coin=invitee_coin,
    )
    db.add(rel)
    # 双方发金币(同事务,与建关系一起 commit)。ref_id 互指对方便于对账。
    crud_wallet.grant_coins(
        db, inviter.id, inviter_coin,
        biz_type="invite_inviter", ref_id=str(invitee.id), remark="邀请好友奖励",
    )
    crud_wallet.grant_coins(
        db, invitee.id, invitee_coin,
        biz_type="invite_invitee", ref_id=str(inviter.id), remark="新人受邀奖励",
    )
    try:
        db.commit()
    except IntegrityError:
        # 并发:同一 invitee 另一个请求先建了关系 → 回滚返回已存在(幂等兜底)
        db.rollback()
        existing = _relation_of_invitee(db, invitee.id)
        if existing is not None:
            return BindResult("already_bound", existing)
        raise
    except Exception:
        # 其它 commit 失败(DB 故障 / PG 序列化冲突等):显式回滚,保证"建关系 + 双方发币"
        # 原子(要么全成要么全无),不依赖 get_db 关闭时的隐式回滚,语义更硬。
        db.rollback()
        raise
    db.refresh(rel)
    return BindResult("success", rel, invitee_coin)


def get_stats(db: Session, inviter_id: int) -> tuple[int, int]:
    """返回 (已成功邀请人数, 累计从邀请获得的金币)。"""
    count = db.execute(
        select(func.count())
        .select_from(InviteRelation)
        .where(InviteRelation.inviter_user_id == inviter_id)
    ).scalar_one()
    coins = db.execute(
        select(func.coalesce(func.sum(InviteRelation.inviter_coin), 0))
        .where(InviteRelation.inviter_user_id == inviter_id)
    ).scalar_one()
    return int(count), int(coins)