meida_backend/packages/domain/billing/quota.py

"""Quota arithmetic — pure logic, no IO (架构文档 §7.4 V1 配额三段式).

The three-phase protocol lives here as data transformations so it can be unit
tested without a database:

    reserve  → 冻结额度并写 bl_ledger(status=reserved)
    commit   → 确认扣除 (status=committed)
    refund   → 释放额度并写反向流水 (tx_type=refund)

V1 doesn't involve money, but 用户对「次数被白扣」的敏感度不比扣钱低，and V1
has no customer-service refund path — the mechanism is the only guarantee.
"""

from __future__ import annotations

from dataclasses import dataclass, replace
from datetime import UTC, datetime, timedelta, timezone
from typing import Literal, get_args

from packages.domain.errors import QuotaExceededError

PeriodType = Literal["day", "month", "lifetime"]
#: 合法周期取值集合，供配置校验复用；单一事实来源是上面的 Literal。
PERIOD_TYPES: frozenset[str] = frozenset(get_args(PeriodType))
LedgerStatus = Literal["reserved", "committed", "refunded"]
TxType = Literal["system_grant", "consume", "refund", "purchase", "gift"]

#: FR-BILL-01: 跨时区按 UTC+8 统一。
DEFAULT_TZ_OFFSET_HOURS = 8


def period_key(
    now: datetime,
    period_type: PeriodType,
    tz_offset_hours: int = DEFAULT_TZ_OFFSET_HOURS,
) -> str:
    """Bucket key that makes quota reset naturally, with no cron job.

    ``day`` → ``d:2026-08-13``; ``month`` → ``m:2026-08``; ``lifetime`` → ``life``.
    """
    if period_type == "lifetime":
        return "life"
    local = now.astimezone(UTC) + timedelta(hours=tz_offset_hours)
    if period_type == "day":
        return f"d:{local:%Y-%m-%d}"
    if period_type == "month":
        return f"m:{local:%Y-%m}"
    msg = f"unknown period_type: {period_type}"
    raise ValueError(msg)


def period_reset_at(
    now: datetime,
    period_type: PeriodType,
    tz_offset_hours: int = DEFAULT_TZ_OFFSET_HOURS,
) -> datetime | None:
    """When the current period rolls over (used for the `reset_at` hint)."""
    if period_type == "lifetime":
        return None
    tz = timezone(timedelta(hours=tz_offset_hours))
    local = now.astimezone(tz)
    if period_type == "day":
        nxt = (local + timedelta(days=1)).replace(hour=0, minute=0, second=0, microsecond=0)
    else:
        year, month = (local.year + 1, 1) if local.month == 12 else (local.year, local.month + 1)
        nxt = local.replace(
            year=year, month=month, day=1, hour=0, minute=0, second=0, microsecond=0
        )
    return nxt.astimezone(UTC)


@dataclass(frozen=True, slots=True)
class QuotaState:
    """Current period counters for one (user, quota_key)."""

    used: int
    base_limit: int
    bonus: int = 0
    hard_cap: int | None = None

    @property
    def effective_limit(self) -> int:
        raw = self.base_limit + self.bonus
        if self.hard_cap is None:
            return raw
        return min(raw, self.hard_cap)

    @property
    def remaining(self) -> int:
        return max(self.effective_limit - self.used, 0)


@dataclass(frozen=True, slots=True)
class QuotaDecision:
    allowed: bool
    required: int
    remaining: int
    effective_limit: int


def evaluate(state: QuotaState, amount: int) -> QuotaDecision:
    """Server-side hard check (FR-BILL-06). Client hints are UX only."""
    if amount <= 0:
        msg = "amount must be positive"
        raise ValueError(msg)
    return QuotaDecision(
        allowed=state.used + amount <= state.effective_limit,
        required=amount,
        remaining=state.remaining,
        effective_limit=state.effective_limit,
    )


def reserve(state: QuotaState, amount: int) -> QuotaState:
    """Freeze ``amount`` units; raises :class:`QuotaExceededError` when over the limit."""
    decision = evaluate(state, amount)
    if not decision.allowed:
        raise QuotaExceededError(required=amount, remaining=decision.remaining)
    return replace(state, used=state.used + amount)


def refund(state: QuotaState, amount: int) -> QuotaState:
    """Release a previously reserved amount. Never drives ``used`` below zero."""
    if amount <= 0:
        msg = "amount must be positive"
        raise ValueError(msg)
    return replace(state, used=max(state.used - amount, 0))


@dataclass(frozen=True, slots=True)
class BonusDecision:
    """任务奖励发放的封顶判定结果（FR-BILL-03）.

    ``allowed=False`` 时 ``bonus/remaining/effective_limit`` 保持发放前的值：
    调用方必须整体放弃，不做部分发放 ——「做任务换 N 次」对用户是明确承诺，
    发一半既说不清也无法用流水对账。
    """

    allowed: bool
    amount: int
    bonus: int
    remaining: int
    effective_limit: int


def grant_bonus(state: QuotaState, amount: int, task_bonus_max: int) -> BonusDecision:
    """把任务奖励 ``amount`` 计入 ``state.bonus``，超上限则整体拒绝（FR-BILL-03）.

    - 只动 ``bonus``，不动 ``used``：任务奖励抬高的是额度，不是用量；
    - ``task_bonus_max`` 来自 ``bl_quota_rule``，是任务奖励独立于 ``hard_cap``
      的第二道天花板（两者都生效时 effective_limit 取更紧的）；
    - 边界语义：恰好到顶（bonus + amount == max）允许，超出一格（== max+1）拒绝。
    """
    if amount <= 0:
        msg = "amount must be positive"
        raise ValueError(msg)
    if state.bonus + amount > task_bonus_max:
        return BonusDecision(
            allowed=False,
            amount=amount,
            bonus=state.bonus,
            remaining=state.remaining,
            effective_limit=state.effective_limit,
        )
    granted = replace(state, bonus=state.bonus + amount)
    return BonusDecision(
        allowed=True,
        amount=amount,
        bonus=granted.bonus,
        remaining=granted.remaining,
        effective_limit=granted.effective_limit,
    )


def balance_after(state: QuotaState) -> int:
    """Value written to ``bl_ledger.balance_after``.

    V1 has no wallet balance, so the ledger records *remaining quota* — that is
    the number the user actually sees ("今日剩余 N 次"). When V2 introduces
    ``bl_wallet`` this becomes the wallet balance without a schema change.
    """
    return state.remaining


def reserve_idempotency_key(*, user_id: int, quota_key: str, request_key: str) -> str:
    return f"rsv:{user_id}:{quota_key}:{request_key}"


def grant_idempotency_key(*, user_id: int, task_key: str, claim_period_key: str) -> str:
    """FR-BILL-03 发放幂等：同一任务同一周期派生同一个键.

    周期取任务的 ``claim_period``（不是额度行的周期）——「每天可领」的月度
    额度任务，幂等粒度必须跟任务承诺走。
    """
    return f"grant:{user_id}:{task_key}:{claim_period_key}"


def refund_idempotency_key(reserve_key: str) -> str:
    """FR-BILL-05: 退还必须幂等。Derived, so a retry cannot double-refund."""
    return f"rfd:{reserve_key}"