xiaoxia-saas/packages/domain/quota.py

"""Quota system with registry pattern.

Member tiers (see packages.domain.points_rules.MEMBERSHIP_PRICES):
- free:      2GB storage, 5 videos/month, 3 concurrent, 3 templates, 10 voiceovers, no AI voice
- monthly:   月卡会员（同 basic 级别）
- quarterly: 季卡会员（同 premium 级别）
- yearly:    年卡会员（同 premium 级别，更多每日免费额度）

旧档位（standard/pro/enterprise/basic/premium）已在 #1894 清理，统一为 free/monthly/quarterly/yearly。
Quota dimensions are registered by modules via the ModuleRegistry,
and checked against the user's membership type.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Optional


class QuotaDimension(str, Enum):
    """配额维度 - 所有可量化的资源限制"""

    STORAGE_GB = "storage_gb"  # 存储空间 (GB)
    VIDEOS_PER_MONTH = "videos_per_month"  # 每月生成视频数
    MAX_CONCURRENT = "max_concurrent"  # 最大并发任务数
    MAX_TEMPLATES = "max_templates"  # 最大模板数
    MAX_TITLES = "max_titles"  # 最大标题库条目数
    MAX_VOICEOVERS = "max_voiceovers"  # 最大配音库条目数
    AI_VOICE_ENABLED = "ai_voice_enabled"  # AI 配音是否可用 (0/1)
    # 以下维度由扩展模块注册，初始配额为 0（由模块注册时填充）
    AI_VOICE_CREDITS = "ai_voice_credits"  # AI 配音积分（每月）
    BATCH_EXPORT_ENABLED = "batch_export_enabled"  # 批量导出
    MULTI_PLATFORM_ENABLED = "multi_platform_enabled"  # 多平台发布
    DEDUP_REPORT_ENABLED = "dedup_report_enabled"  # 去重检测报告


@dataclass
class QuotaTier:
    """一个套餐等级的配额定义"""

    name: str
    limits: dict[str, float] = field(default_factory=dict)

    def get_limit(self, dimension: str) -> float:
        """获取指定维度的配额限制，未定义返回 0"""
        return self.limits.get(dimension, 0)

    def is_unlimited(self, dimension: str) -> bool:
        """指定维度是否不限量"""
        return self.limits.get(dimension, float("inf")) == float("inf")


# 三个套餐等级的配额定义
QUOTA_TIERS: dict[str, QuotaTier] = {
    "free": QuotaTier(
        name="free",
        limits={
            QuotaDimension.STORAGE_GB: 2,
            QuotaDimension.VIDEOS_PER_MONTH: 5,
            QuotaDimension.MAX_CONCURRENT: 3,
            QuotaDimension.MAX_TEMPLATES: 3,
            QuotaDimension.MAX_VOICEOVERS: 10,
            QuotaDimension.AI_VOICE_ENABLED: 0,
            QuotaDimension.AI_VOICE_CREDITS: 0,
            QuotaDimension.BATCH_EXPORT_ENABLED: 0,
            QuotaDimension.MULTI_PLATFORM_ENABLED: 0,
            QuotaDimension.DEDUP_REPORT_ENABLED: 0,
        },
    ),
    # 月卡会员：基础付费档（原 basic）
    "monthly": QuotaTier(
        name="monthly",
        limits={
            QuotaDimension.STORAGE_GB: 20,
            QuotaDimension.VIDEOS_PER_MONTH: 30,
            QuotaDimension.MAX_CONCURRENT: 10,
            QuotaDimension.MAX_TEMPLATES: 15,
            QuotaDimension.MAX_VOICEOVERS: 100,
            QuotaDimension.AI_VOICE_ENABLED: 1,
            QuotaDimension.AI_VOICE_CREDITS: 100,
            QuotaDimension.BATCH_EXPORT_ENABLED: 1,
            QuotaDimension.MULTI_PLATFORM_ENABLED: 0,
            QuotaDimension.DEDUP_REPORT_ENABLED: 0,
        },
    ),
    # 季卡会员：高级付费档（原 premium）
    "quarterly": QuotaTier(
        name="quarterly",
        limits={
            QuotaDimension.STORAGE_GB: 100,
            QuotaDimension.VIDEOS_PER_MONTH: 100,
            QuotaDimension.MAX_CONCURRENT: 20,
            QuotaDimension.MAX_TEMPLATES: float("inf"),
            QuotaDimension.MAX_VOICEOVERS: 100,
            QuotaDimension.AI_VOICE_ENABLED: 1,
            QuotaDimension.AI_VOICE_CREDITS: 500,
            QuotaDimension.BATCH_EXPORT_ENABLED: 1,
            QuotaDimension.MULTI_PLATFORM_ENABLED: 1,
            QuotaDimension.DEDUP_REPORT_ENABLED: 1,
        },
    ),
    # 年卡会员：同季卡配额 + 每日不限免费条数（由前端/积分规则实现）
    "yearly": QuotaTier(
        name="yearly",
        limits={
            QuotaDimension.STORAGE_GB: 100,
            QuotaDimension.VIDEOS_PER_MONTH: float("inf"),
            QuotaDimension.MAX_CONCURRENT: 20,
            QuotaDimension.MAX_TEMPLATES: float("inf"),
            QuotaDimension.MAX_VOICEOVERS: 200,
            QuotaDimension.AI_VOICE_ENABLED: 1,
            QuotaDimension.AI_VOICE_CREDITS: 2000,
            QuotaDimension.BATCH_EXPORT_ENABLED: 1,
            QuotaDimension.MULTI_PLATFORM_ENABLED: 1,
            QuotaDimension.DEDUP_REPORT_ENABLED: 1,
        },
    ),
}

# #1894: 旧档位别名（basic/standard → monthly, premium/pro/enterprise → quarterly）
# 历史 DB 数据、单测和内部模块可能仍在传旧 plan_name；这里保留别名保证配额查询不炸。
# 新代码请统一使用 free/monthly/quarterly/yearly。
QUOTA_TIERS["basic"] = QUOTA_TIERS["monthly"]
QUOTA_TIERS["standard"] = QUOTA_TIERS["monthly"]
QUOTA_TIERS["premium"] = QUOTA_TIERS["quarterly"]
QUOTA_TIERS["pro"] = QUOTA_TIERS["quarterly"]
QUOTA_TIERS["enterprise"] = QUOTA_TIERS["quarterly"]


class QuotaWarningLevel:
    """配额告警级别"""

    NORMAL = "normal"  # 使用量 < 80%
    WARNING = "warning"  # 80% <= 使用量 < 100%
    CRITICAL = "critical"  # 95% <= 使用量 < 100%
    EXCEEDED = "exceeded"  # 使用量 >= 100%


@dataclass
class QuotaCheckResult:
    """配额检查结果"""

    allowed: bool
    dimension: str
    limit: float
    used: float
    remaining: float
    warning_level: str

    @property
    def usage_percent(self) -> float:
        if self.limit <= 0:
            return 100.0 if self.used > 0 else 0.0
        if self.limit == float("inf"):
            return 0.0
        return min(100.0, (self.used / self.limit) * 100)


class QuotaRegistry:
    """配额注册表

    管理所有配额维度的定义和套餐限制。
    扩展模块可通过 register_dimension() 注册新的配额维度。
    """

    def __init__(self) -> None:
        self._dimensions: dict[str, str] = {}  # dimension_name -> description
        self._tiers: dict[str, QuotaTier] = dict(QUOTA_TIERS)

        # 初始化内置维度
        for dim in QuotaDimension:
            self._dimensions[dim.value] = dim.name

    def register_dimension(
        self,
        dimension: str,
        description: str,
        default_limits: Optional[dict[str, float]] = None,
    ) -> None:
        """注册新的配额维度

        Args:
            dimension: 维度名称
            description: 人类可读描述
            default_limits: 各套餐的默认限制 {plan_name: limit}
        """
        if dimension in self._dimensions:
            return  # 幂等

        self._dimensions[dimension] = description

        # 为各套餐设置默认限制
        if default_limits:
            for plan_name, limit in default_limits.items():
                if plan_name in self._tiers:
                    self._tiers[plan_name].limits[dimension] = limit
        else:
            # 默认所有套餐该维度为 0
            for tier in self._tiers.values():
                tier.limits.setdefault(dimension, 0)

    def get_tier(self, plan_name: str) -> Optional[QuotaTier]:
        """获取套餐配额定义"""
        return self._tiers.get(plan_name)

    def get_limit(self, plan_name: str, dimension: str) -> float:
        """获取指定套餐在指定维度的限制"""
        tier = self._tiers.get(plan_name)
        if tier is None:
            return 0
        return tier.get_limit(dimension)

    def list_dimensions(self) -> dict[str, str]:
        """列出所有已注册的配额维度"""
        return dict(self._dimensions)

    def list_tiers(self) -> list[str]:
        """列出所有套餐等级"""
        return list(self._tiers.keys())


class QuotaChecker:
    """配额检查器

    检查用户在指定维度的使用量是否超出限制。
    需要调用方提供当前使用量（used），本类不直接访问数据库。
    """

    def __init__(self, registry: Optional[QuotaRegistry] = None) -> None:
        self._registry = registry or QuotaRegistry()

    def check(
        self,
        plan_name: str,
        dimension: str,
        used: float,
    ) -> QuotaCheckResult:
        """检查指定维度的配额使用情况

        Args:
            plan_name: 会员类型 (free/monthly/quarterly/yearly)
            dimension: 配额维度
            used: 当前已使用量

        Returns:
            QuotaCheckResult
        """
        limit = self._registry.get_limit(plan_name, dimension)
        remaining = max(0, limit - used) if limit != float("inf") else float("inf")
        allowed = used < limit if limit != float("inf") else True
        warning_level = self._compute_warning_level(used, limit)

        return QuotaCheckResult(
            allowed=allowed,
            dimension=dimension,
            limit=limit,
            used=used,
            remaining=remaining,
            warning_level=warning_level,
        )

    def check_multiple(
        self,
        plan_name: str,
        usage: dict[str, float],
    ) -> list[QuotaCheckResult]:
        """批量检查多个维度的配额"""
        return [self.check(plan_name, dim, used) for dim, used in usage.items()]

    @staticmethod
    def _compute_warning_level(used: float, limit: float) -> str:
        """计算告警级别"""
        if limit <= 0:
            return QuotaWarningLevel.EXCEEDED if used > 0 else QuotaWarningLevel.NORMAL
        if limit == float("inf"):
            return QuotaWarningLevel.NORMAL

        ratio = used / limit
        if ratio >= 1.0:
            return QuotaWarningLevel.EXCEEDED
        if ratio >= 0.95:
            return QuotaWarningLevel.CRITICAL
        if ratio >= 0.80:
            return QuotaWarningLevel.WARNING
        return QuotaWarningLevel.NORMAL


def get_warning_level(used: float, limit: float) -> str:
    """便捷函数：计算配额告警级别"""
    return QuotaChecker._compute_warning_level(used, limit)


# 全局单例
quota_registry = QuotaRegistry()
quota_checker = QuotaChecker(quota_registry)