xiaoxia-saas/packages/domain/quota.py

"""
配额检查服务
用于检查工作空间是否超出配额限制
"""
from typing import Optional


class QuotaChecker:
    """配额检查器"""
    
    def __init__(
        self,
        workspace_repository,
        project_repository,
    ):
        self.workspace_repository = workspace_repository
        self.project_repository = project_repository
    
    def check_can_create_project(
        self,
        workspace_id: str,
    ) -> tuple[bool, Optional[str]]:
        """
        检查是否可以创建项目
        
        Args:
            workspace_id: 工作空间 ID
        
        Returns:
            (是否可以, 错误信息)
        """
        workspace = self.workspace_repository.find_by_id(workspace_id)
        if not workspace:
            return False, "Workspace not found"
        
        # 获取当前项目数量
        current_count = self.project_repository.count_by_workspace(workspace_id)
        
        # 检查是否超出配额（999999 表示无限）
        if workspace.max_projects != 999999 and current_count >= workspace.max_projects:
            return False, f"Project limit reached ({workspace.max_projects}). Upgrade your plan to create more projects."
        
        return True, None
    
    def check_storage_available(
        self,
        workspace_id: str,
        additional_gb: float,
    ) -> tuple[bool, Optional[str]]:
        """
        检查存储空间是否足够
        
        Args:
            workspace_id: 工作空间 ID
            additional_gb: 需要的额外存储空间（GB）
        
        Returns:
            (是否可以, 错误信息)
        """
        workspace = self.workspace_repository.find_by_id(workspace_id)
        if not workspace:
            return False, "Workspace not found"
        
        # 检查存储空间
        new_usage = workspace.used_storage_gb + additional_gb
        
        if new_usage > workspace.max_storage_gb:
            remaining = workspace.max_storage_gb - workspace.used_storage_gb
            return False, f"Storage limit exceeded. Available: {remaining:.2f}GB, Required: {additional_gb:.2f}GB. Upgrade your plan for more storage."
        
        return True, None
    
    def get_quota_status(self, workspace_id: str) -> dict:
        """
        获取配额使用状态
        
        Args:
            workspace_id: 工作空间 ID
        
        Returns:
            配额状态信息
        """
        workspace = self.workspace_repository.find_by_id(workspace_id)
        if not workspace:
            return None
        
        # 获取项目数量
        project_count = self.project_repository.count_by_workspace(workspace_id)
        
        # 计算使用率
        project_usage_percent = (
            (project_count / workspace.max_projects * 100)
            if workspace.max_projects != 999999
            else 0  # 无限制
        )
        
        storage_usage_percent = (
            (workspace.used_storage_gb / workspace.max_storage_gb * 100)
            if workspace.max_storage_gb > 0
            else 0
        )
        
        return {
            "workspace_id": workspace.id,
            "subscription_plan": workspace.subscription_plan,
            "projects": {
                "used": project_count,
                "limit": workspace.max_projects,
                "unlimited": workspace.max_projects == 999999,
                "usage_percent": project_usage_percent,
            },
            "storage": {
                "used_gb": workspace.used_storage_gb,
                "limit_gb": workspace.max_storage_gb,
                "remaining_gb": workspace.max_storage_gb - workspace.used_storage_gb,
                "usage_percent": storage_usage_percent,
            },
        }
    
    def update_storage_usage(
        self,
        workspace_id: str,
        delta_gb: float,
    ) -> tuple[bool, Optional[str]]:
        """
        更新存储使用量
        
        Args:
            workspace_id: 工作空间 ID
            delta_gb: 变化量（正数为增加，负数为减少）
        
        Returns:
            (是否成功, 错误信息)
        """
        workspace = self.workspace_repository.find_by_id(workspace_id)
        if not workspace:
            return False, "Workspace not found"
        
        # 更新使用量
        new_usage = workspace.used_storage_gb + delta_gb
        
        # 不能为负数
        if new_usage < 0:
            new_usage = 0
        
        workspace.used_storage_gb = new_usage
        self.workspace_repository.save(workspace)
        
        return True, None


class QuotaWarningLevel:
    """配额警告级别"""
    
    NORMAL = "normal"  # <80%
    WARNING = "warning"  # 80-90%
    CRITICAL = "critical"  # 90-100%
    EXCEEDED = "exceeded"  # >100%


def get_warning_level(usage_percent: float) -> str:
    """
    根据使用率获取警告级别
    
    Args:
        usage_percent: 使用率（0-100）
    
    Returns:
        警告级别
    """
    if usage_percent >= 100:
        return QuotaWarningLevel.EXCEEDED
    elif usage_percent >= 90:
        return QuotaWarningLevel.CRITICAL
    elif usage_percent >= 80:
        return QuotaWarningLevel.WARNING
    else:
        return QuotaWarningLevel.NORMAL