Aether/src/services/user/apikey.py

"""
API密钥管理服务
"""

from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional

from sqlalchemy import func
from sqlalchemy.orm import Session

from src.core.crypto import crypto_service
from src.core.logger import logger
from src.models.database import ApiKey, Usage, User



class ApiKeyService:
    """API密钥管理服务"""

    @staticmethod
    def create_api_key(
        db: Session,
        user_id: str,  # UUID
        name: Optional[str] = None,
        allowed_providers: Optional[List[str]] = None,
        allowed_api_formats: Optional[List[str]] = None,
        allowed_models: Optional[List[str]] = None,
        rate_limit: int = 100,
        concurrent_limit: int = 5,
        expire_days: Optional[int] = None,
        initial_balance_usd: Optional[float] = None,
        is_standalone: bool = False,
        auto_delete_on_expiry: bool = False,
    ) -> tuple[ApiKey, str]:
        """创建新的API密钥，返回密钥对象和明文密钥

        Args:
            db: 数据库会话
            user_id: 用户ID
            name: 密钥名称
            allowed_providers: 允许的提供商列表
            allowed_api_formats: 允许的 API 格式列表
            allowed_models: 允许的模型列表
            rate_limit: 速率限制
            concurrent_limit: 并发限制
            expire_days: 过期天数，None = 永不过期
            initial_balance_usd: 初始余额（USD），仅用于独立Key，None = 无限制
            is_standalone: 是否为独立余额Key（仅管理员可创建）
            auto_delete_on_expiry: 过期后是否自动删除（True=物理删除，False=仅禁用）
        """

        # 生成密钥
        key = ApiKey.generate_key()
        key_hash = ApiKey.hash_key(key)
        key_encrypted = crypto_service.encrypt(key)  # 加密存储密钥

        # 计算过期时间
        expires_at = None
        if expire_days:
            expires_at = datetime.now(timezone.utc) + timedelta(days=expire_days)

        # 空数组转为 None（表示不限制）
        api_key = ApiKey(
            user_id=user_id,
            key_hash=key_hash,
            key_encrypted=key_encrypted,
            name=name or f"API Key {datetime.now(timezone.utc).strftime('%Y%m%d%H%M%S')}",
            allowed_providers=allowed_providers or None,
            allowed_api_formats=allowed_api_formats or None,
            allowed_models=allowed_models or None,
            rate_limit=rate_limit,
            concurrent_limit=concurrent_limit,
            expires_at=expires_at,
            balance_used_usd=0.0,
            current_balance_usd=initial_balance_usd,  # 直接使用初始余额，None = 无限制
            is_standalone=is_standalone,
            auto_delete_on_expiry=auto_delete_on_expiry,
            is_active=True,
        )

        db.add(api_key)
        db.commit()
        db.refresh(api_key)

        logger.info(f"创建API密钥: 用户ID {user_id}, 密钥名 {api_key.name}, "
            f"独立Key={is_standalone}, 初始余额={initial_balance_usd}")
        return api_key, key  # 返回密钥对象和明文密钥

    @staticmethod
    def get_api_key(db: Session, key_id: str) -> Optional[ApiKey]:  # UUID
        """获取API密钥"""
        return db.query(ApiKey).filter(ApiKey.id == key_id).first()

    @staticmethod
    def get_api_key_by_key(db: Session, key: str) -> Optional[ApiKey]:
        """通过密钥字符串获取API密钥"""
        key_hash = ApiKey.hash_key(key)
        return db.query(ApiKey).filter(ApiKey.key_hash == key_hash).first()

    @staticmethod
    def list_user_api_keys(
        db: Session, user_id: str, is_active: Optional[bool] = None  # UUID
    ) -> List[ApiKey]:
        """列出用户的所有API密钥（不包括独立Key）"""
        query = db.query(ApiKey).filter(
            ApiKey.user_id == user_id, ApiKey.is_standalone == False  # 排除独立Key
        )

        if is_active is not None:
            query = query.filter(ApiKey.is_active == is_active)

        return query.order_by(ApiKey.created_at.desc()).all()

    @staticmethod
    def list_standalone_api_keys(db: Session, is_active: Optional[bool] = None) -> List[ApiKey]:
        """列出所有独立余额Key（仅管理员可用）"""
        query = db.query(ApiKey).filter(ApiKey.is_standalone == True)

        if is_active is not None:
            query = query.filter(ApiKey.is_active == is_active)

        return query.order_by(ApiKey.created_at.desc()).all()

    @staticmethod
    def update_api_key(db: Session, key_id: str, **kwargs) -> Optional[ApiKey]:  # UUID
        """更新API密钥"""
        api_key = db.query(ApiKey).filter(ApiKey.id == key_id).first()
        if not api_key:
            return None

        # 可更新的字段
        updatable_fields = [
            "name",
            "allowed_providers",
            "allowed_api_formats",
            "allowed_models",
            "rate_limit",
            "concurrent_limit",
            "is_active",
            "expires_at",
            "balance_limit_usd",
            "auto_delete_on_expiry",
        ]

        # 允许显式设置为空数组/None 的字段（空数组会转为 None，表示"全部"）
        nullable_list_fields = {"allowed_providers", "allowed_api_formats", "allowed_models"}

        for field, value in kwargs.items():
            if field not in updatable_fields:
                continue
            # 对于 nullable_list_fields，空数组应该转为 None（表示不限制）
            if field in nullable_list_fields:
                if value is not None:
                    # 空数组转为 None（表示允许全部）
                    setattr(api_key, field, value if value else None)
            elif value is not None:
                setattr(api_key, field, value)

        api_key.updated_at = datetime.now(timezone.utc)
        db.commit()
        db.refresh(api_key)

        logger.debug(f"更新API密钥: ID {key_id}")
        return api_key

    @staticmethod
    def delete_api_key(db: Session, key_id: str) -> bool:  # UUID
        """删除API密钥（禁用）"""
        api_key = db.query(ApiKey).filter(ApiKey.id == key_id).first()
        if not api_key:
            return False

        api_key.is_active = False
        api_key.updated_at = datetime.now(timezone.utc)
        db.commit()

        logger.info(f"删除API密钥: ID {key_id}")
        return True

    @staticmethod
    def get_remaining_balance(api_key: ApiKey) -> Optional[float]:
        """计算剩余余额（仅用于独立Key）

        Returns:
            剩余余额，None 表示无限制或非独立Key
        """
        if not api_key.is_standalone:
            return None

        if api_key.current_balance_usd is None:
            return None

        # 剩余余额 = 当前余额 - 已使用余额
        remaining = api_key.current_balance_usd - (api_key.balance_used_usd or 0)
        return max(0, remaining)  # 不能为负数

    @staticmethod
    def check_balance(api_key: ApiKey) -> tuple[bool, Optional[float]]:
        """检查余额限制（仅用于独立Key）

        Returns:
            (is_allowed, remaining_balance): 是否允许请求，剩余余额（None表示无限制）
        """
        if not api_key.is_standalone:
            # 非独立Key不检查余额
            return True, None

        # 使用新的预付费模式: current_balance_usd
        if api_key.current_balance_usd is None:
            # 无余额限制
            return True, None

        # 使用统一的余额计算方法
        remaining = ApiKeyService.get_remaining_balance(api_key)
        is_allowed = remaining > 0 if remaining is not None else True

        if not is_allowed:
            logger.warning(f"API密钥余额不足: Key ID {api_key.id}, " f"剩余余额 ${remaining:.4f}")

        return is_allowed, remaining

    @staticmethod
    def check_rate_limit(db: Session, api_key: ApiKey, window_minutes: int = 1) -> tuple[bool, int]:
        """检查速率限制

        Returns:
            (is_allowed, remaining): 是否允许请求，剩余可用次数
            当 rate_limit 为 None 时表示不限制，返回 (True, -1)
        """
        # 如果 rate_limit 为 None，表示不限制
        if api_key.rate_limit is None:
            return True, -1  # -1 表示无限制

        # 计算时间窗口
        window_start = datetime.now(timezone.utc) - timedelta(minutes=window_minutes)

        # 统计窗口内的请求数
        request_count = (
            db.query(func.count(Usage.id))
            .filter(Usage.api_key_id == api_key.id, Usage.created_at >= window_start)
            .scalar()
            or 0
        )

        # 检查是否超限
        is_allowed = request_count < api_key.rate_limit

        if not is_allowed:
            logger.warning(f"API密钥速率限制: Key ID {api_key.id}, 请求数 {request_count}/{api_key.rate_limit}")

        return is_allowed, api_key.rate_limit - request_count

    @staticmethod
    def add_balance(db: Session, key_id: str, amount_usd: float) -> Optional[ApiKey]:
        """为独立余额Key调整余额

        Args:
            db: 数据库会话
            key_id: API Key ID
            amount_usd: 要调整的余额金额（USD），正数为增加，负数为扣除

        Returns:
            更新后的API Key对象，如果Key不存在或不是独立Key则返回None
        """
        api_key = db.query(ApiKey).filter(ApiKey.id == key_id).first()
        if not api_key:
            logger.warning(f"余额调整失败: Key ID {key_id} 不存在")
            return None

        if not api_key.is_standalone:
            logger.warning(f"余额调整失败: Key ID {key_id} 不是独立余额Key")
            return None

        if amount_usd == 0:
            logger.warning(f"余额调整失败: 调整金额不能为0，当前值 ${amount_usd}")
            return None

        # 如果是扣除（负数），检查是否超过当前余额
        if amount_usd < 0:
            current = api_key.current_balance_usd or 0
            if abs(amount_usd) > current:
                logger.warning(f"余额扣除失败: 扣除金额 ${abs(amount_usd):.4f} 超过当前余额 ${current:.4f}")
                return None

        # 调整当前余额
        if api_key.current_balance_usd is None:
            api_key.current_balance_usd = amount_usd if amount_usd > 0 else 0
        else:
            api_key.current_balance_usd = max(0, api_key.current_balance_usd + amount_usd)

        api_key.updated_at = datetime.now(timezone.utc)
        db.commit()
        db.refresh(api_key)

        action = "增加" if amount_usd > 0 else "扣除"
        logger.info(f"余额调整成功: Key ID {key_id}, {action} ${abs(amount_usd):.4f}, "
            f"新余额 ${api_key.current_balance_usd:.4f}")
        return api_key

    @staticmethod
    def cleanup_expired_keys(db: Session, auto_delete: bool = False) -> int:
        """清理过期的API密钥

        Args:
            db: 数据库会话
            auto_delete: 全局默认行为（True=物理删除，False=仅禁用）
                        单个Key的 auto_delete_on_expiry 字段会覆盖此设置

        Returns:
            int: 清理的密钥数量
        """
        now = datetime.now(timezone.utc)
        expired_keys = (
            db.query(ApiKey)
            .filter(ApiKey.expires_at <= now, ApiKey.is_active == True)  # 只处理仍然活跃的
            .all()
        )

        count = 0
        for api_key in expired_keys:
            # 优先使用Key自身的auto_delete_on_expiry设置,否则使用全局设置
            should_delete = (
                api_key.auto_delete_on_expiry
                if api_key.auto_delete_on_expiry is not None
                else auto_delete
            )

            if should_delete:
                # 物理删除（Usage记录会保留，因为是 SET NULL）
                db.delete(api_key)
                logger.info(f"删除过期API密钥: ID {api_key.id}, 名称 {api_key.name}, "
                    f"过期时间 {api_key.expires_at}")
            else:
                # 仅禁用
                api_key.is_active = False
                api_key.updated_at = now
                logger.info(f"禁用过期API密钥: ID {api_key.id}, 名称 {api_key.name}, "
                    f"过期时间 {api_key.expires_at}")
            count += 1

        if count > 0:
            db.commit()

        return count

    @staticmethod
    def get_api_key_stats(
        db: Session,
        key_id: str,  # UUID
        start_date: Optional[datetime] = None,
        end_date: Optional[datetime] = None,
    ) -> Dict[str, Any]:
        """获取API密钥使用统计"""

        api_key = db.query(ApiKey).filter(ApiKey.id == key_id).first()
        if not api_key:
            return {}

        query = db.query(Usage).filter(Usage.api_key_id == key_id)

        if start_date:
            query = query.filter(Usage.created_at >= start_date)
        if end_date:
            query = query.filter(Usage.created_at <= end_date)

        # 统计数据
        stats = db.query(
            func.count(Usage.id).label("requests"),
            func.sum(Usage.total_tokens).label("tokens"),
            func.sum(Usage.total_cost_usd).label("cost_usd"),
            func.avg(Usage.response_time_ms).label("avg_response_time"),
        ).filter(Usage.api_key_id == key_id)

        if start_date:
            stats = stats.filter(Usage.created_at >= start_date)
        if end_date:
            stats = stats.filter(Usage.created_at <= end_date)

        result = stats.first()

        # 按天统计
        daily_stats = db.query(
            func.date(Usage.created_at).label("date"),
            func.count(Usage.id).label("requests"),
            func.sum(Usage.total_tokens).label("tokens"),
            func.sum(Usage.total_cost_usd).label("cost_usd"),
        ).filter(Usage.api_key_id == key_id)

        if start_date:
            daily_stats = daily_stats.filter(Usage.created_at >= start_date)
        if end_date:
            daily_stats = daily_stats.filter(Usage.created_at <= end_date)

        daily_stats = daily_stats.group_by(func.date(Usage.created_at)).all()

        return {
            "key_id": key_id,
            "key_name": api_key.name,
            "total_requests": result.requests or 0,
            "total_tokens": result.tokens or 0,
            "total_cost_usd": float(result.cost_usd or 0),
            "avg_response_time_ms": float(result.avg_response_time or 0),
            "daily_stats": [
                {
                    "date": stat.date.isoformat() if stat.date else None,
                    "requests": stat.requests,
                    "tokens": stat.tokens,
                    "cost_usd": float(stat.cost_usd),
                }
                for stat in daily_stats
            ],
        }