yingqing/core/backend/apps/billing/services/ledger.py

from decimal import Decimal

from django.db import transaction
from django.db.models import Sum
from django.utils import timezone

from apps.billing.models import CreditAccount, CreditLedger, CreditReservation, QuotaPolicy


def _enforce_member_monthly_limit(*, team, user, amount: Decimal) -> None:
    """成员月度限额管控。团队页可设 monthly_credit_limit(0=不限),此前只存不管——纯装饰。
    口径与成员列表的 month_charged 一致:自然月内该成员的 CHARGE 流水合计;
    再加上其当前 ACTIVE 预留(在途任务),防止并发把限额冲穿。
    调用方(reserve_credit)已持有 account 行锁,团队内 reserve 串行化,此检查无竞态。"""
    if user is None:
        return
    member = team.members.filter(user=user).first()
    if member is None:
        return
    limit = member.monthly_credit_limit or Decimal("0")
    if limit <= 0:
        return  # 0 = 不限额
    now = timezone.now()
    month_start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
    charged = (
        CreditLedger.objects.filter(team=team, user=user, ledger_type=CreditLedger.Type.CHARGE, created_at__gte=month_start)
        .aggregate(s=Sum("amount"))["s"] or Decimal("0")
    )
    reserved = (
        CreditReservation.objects.filter(team=team, user=user, status=CreditReservation.Status.ACTIVE)
        .aggregate(s=Sum("amount"))["s"] or Decimal("0")
    )
    if charged + reserved + amount > limit:
        raise ValueError(f"成员本月额度不足:限额 ¥{limit},本月已用 ¥{charged}(另有在途 ¥{reserved})")


def _enforce_quota_policy(*, team, project, amount: Decimal) -> None:
    """平台 4 层额度策略(QuotaPolicy)。**仅当团队存在 active 团队级策略时才生效** ——
    无策略的团队行为完全不变(零回归)。检查 单任务 / 团队月度 / 项目 三类限额(任一为 None 即该维度不限)。"""
    policy = (
        team.quota_policies.filter(is_active=True, user__isnull=True, project__isnull=True)
        .order_by("-created_at")
        .first()
    )
    if policy is None:
        return
    if policy.per_task_limit is not None and amount > policy.per_task_limit:
        raise ValueError(f"单任务额度超限:上限 ¥{policy.per_task_limit}")
    now = timezone.now()
    month_start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
    if policy.monthly_limit is not None:
        charged = (
            CreditLedger.objects.filter(team=team, ledger_type=CreditLedger.Type.CHARGE, created_at__gte=month_start)
            .aggregate(s=Sum("amount"))["s"] or Decimal("0")
        )
        reserved = (
            CreditReservation.objects.filter(team=team, status=CreditReservation.Status.ACTIVE)
            .aggregate(s=Sum("amount"))["s"] or Decimal("0")
        )
        if charged + reserved + amount > policy.monthly_limit:
            raise ValueError(f"团队本月额度超限:上限 ¥{policy.monthly_limit}")
    if policy.project_limit is not None and project is not None:
        p_charged = (
            CreditLedger.objects.filter(team=team, project=project, ledger_type=CreditLedger.Type.CHARGE)
            .aggregate(s=Sum("amount"))["s"] or Decimal("0")
        )
        p_reserved = (
            CreditReservation.objects.filter(team=team, project=project, status=CreditReservation.Status.ACTIVE)
            .aggregate(s=Sum("amount"))["s"] or Decimal("0")
        )
        if p_charged + p_reserved + amount > policy.project_limit:
            raise ValueError(f"项目额度超限:上限 ¥{policy.project_limit}")


@transaction.atomic
def adjust_credit(*, team, amount: Decimal, reason: str = "", operator=None) -> CreditLedger:
    """平台超管手动调额(争议 / 补偿)。amount 可正可负;落 ADJUSTMENT 流水并更新余额。
    调额后余额不得为负(否则拒绝)。"""
    account, _ = CreditAccount.objects.select_for_update().get_or_create(team=team)
    new_balance = account.balance + amount
    if new_balance < 0:
        raise ValueError("调额后余额为负,已拒绝")
    account.balance = new_balance
    account.save(update_fields=["balance", "updated_at"])
    return CreditLedger.objects.create(
        team=team,
        user=operator,
        ledger_type=CreditLedger.Type.ADJUSTMENT,
        amount=amount,
        balance_after=new_balance,
        reason=reason or "平台手动调额",
        metadata={"kind": "admin_adjust"},
    )


@transaction.atomic
def reserve_credit(*, team, user, task, amount: Decimal) -> CreditReservation:
    account, _ = CreditAccount.objects.select_for_update().get_or_create(team=team)
    available = account.balance - account.reserved_balance
    if available < amount:
        raise ValueError("insufficient credit")
    _enforce_member_monthly_limit(team=team, user=user, amount=amount)
    _enforce_quota_policy(team=team, project=task.project, amount=amount)

    account.reserved_balance += amount
    account.save(update_fields=["reserved_balance", "updated_at"])
    reservation = CreditReservation.objects.create(
        team=team,
        user=user,
        project=task.project,
        task=task,
        amount=amount,
    )
    CreditLedger.objects.create(
        team=team,
        user=user,
        project=task.project,
        task=task,
        ledger_type=CreditLedger.Type.RESERVE,
        amount=amount,
        balance_after=account.balance,
        reason="AI 任务预扣额度",
    )
    return reservation


@transaction.atomic
def release_credit(*, reservation: CreditReservation, reason: str = "") -> None:
    account = CreditAccount.objects.select_for_update().get(team=reservation.team)
    # 锁内重读:与 charge 同款竞态——并发双 release(或 charge+release 交错)用陈旧快照都看到 ACTIVE,
    # reserved_balance 会被减两次直接记坏账。已终态(已扣/已释放)→ 幂等返回。
    reservation = CreditReservation.objects.select_for_update().get(id=reservation.id)
    if reservation.status != CreditReservation.Status.ACTIVE:
        return

    account.reserved_balance -= reservation.amount
    account.save(update_fields=["reserved_balance", "updated_at"])
    reservation.status = CreditReservation.Status.RELEASED
    reservation.save(update_fields=["status", "updated_at"])
    CreditLedger.objects.create(
        team=reservation.team,
        user=reservation.user,
        project=reservation.project,
        task=reservation.task,
        ledger_type=CreditLedger.Type.RELEASE,
        amount=reservation.amount,
        balance_after=account.balance,
        reason=reason or "释放预留额度",
    )


@transaction.atomic
def charge_reserved_credit(*, reservation: CreditReservation, actual_amount: Decimal) -> None:
    account = CreditAccount.objects.select_for_update().get(team=reservation.team)
    # 锁内重读 reservation:调用方传入的可能是并发竞态下的陈旧快照(双方都看到 ACTIVE),
    # 旧实现会对同一笔预留扣两次费(实测同秒两条 charge 流水)。已扣费 → 幂等直接返回。
    reservation = CreditReservation.objects.select_for_update().get(id=reservation.id)
    if reservation.status == CreditReservation.Status.CHARGED:
        return
    if reservation.status != CreditReservation.Status.ACTIVE:
        raise ValueError("reservation is not active")
    if actual_amount > reservation.amount:
        raise ValueError("actual amount exceeds reserved amount")

    account.balance -= actual_amount
    account.reserved_balance -= reservation.amount
    account.save(update_fields=["balance", "reserved_balance", "updated_at"])
    reservation.status = CreditReservation.Status.CHARGED
    reservation.save(update_fields=["status", "updated_at"])
    CreditLedger.objects.create(
        team=reservation.team,
        user=reservation.user,
        project=reservation.project,
        task=reservation.task,
        ledger_type=CreditLedger.Type.CHARGE,
        amount=actual_amount,
        balance_after=account.balance,
        reason="AI 任务扣费",
    )
    if reservation.amount > actual_amount:
        CreditLedger.objects.create(
            team=reservation.team,
            user=reservation.user,
            project=reservation.project,
            task=reservation.task,
            ledger_type=CreditLedger.Type.RELEASE,
            amount=reservation.amount - actual_amount,
            balance_after=account.balance,
            reason="释放未用预留额度",
        )