Aether/_deprecated_py_src/services/provider/adapters/kiro/usage.py

"""Kiro usage/quota fetching utilities."""

from __future__ import annotations

import re
import time
import uuid
from typing import Any

import httpx

from src.clients.http_client import HTTPClientPool
from src.core.logger import logger
from src.services.provider.adapters.kiro.headers import (
    build_user_agent_usage,
    build_x_amz_user_agent_usage,
)
from src.services.provider.adapters.kiro.models.credentials import KiroAuthConfig
from src.services.provider.adapters.kiro.models.usage_limits import (
    UsageLimitsResponse,
    calculate_current_usage,
    calculate_total_usage_limit,
)
from src.services.provider.adapters.kiro.request import get_profile_arn_for_payload
from src.services.provider.adapters.kiro.rust_http import execute_kiro_rust_http_request
from src.services.provider.adapters.kiro.token_manager import (
    generate_machine_id,
    is_token_expired,
    refresh_access_token,
)


class KiroAccountBannedException(Exception):
    """Kiro 账户被封禁异常"""

    def __init__(
        self,
        message: str = "账户已封禁",
        status_code: int = 403,
        reason: str | None = None,
    ):
        super().__init__(message)
        self.message = message
        self.status_code = status_code
        self.reason = reason


async def fetch_kiro_usage_limits(
    auth_config: dict[str, Any],
    proxy_config: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """
    调用 Kiro getUsageLimits API 获取使用额度信息

    Args:
        auth_config: 解密后的 KiroAuthConfig 数据
        proxy_config: 代理配置（可选）

    Returns:
        包含 usage_data 和 updated_auth_config 的字典

    Raises:
        RuntimeError: 请求失败时抛出
    """
    cfg = KiroAuthConfig.from_dict(auth_config)

    # 检查是否有缓存的 access_token 且未过期
    access_token: str | None = None
    updated_cfg: KiroAuthConfig | None = None

    if cfg.access_token and not is_token_expired(cfg.expires_at):
        # 使用缓存的 token
        access_token = cfg.access_token
        updated_cfg = cfg
        logger.debug("[KIRO_QUOTA] 使用缓存的 access_token")
    else:
        # token 过期或不存在，需要刷新
        logger.debug("[KIRO_QUOTA] Token 已过期或不存在，正在刷新...")
        access_token, updated_cfg = await refresh_access_token(cfg, proxy_config=proxy_config)

    if not access_token:
        raise RuntimeError("无法获取 Kiro access_token")

    # 构建请求
    effective_cfg = updated_cfg or cfg
    region = effective_cfg.effective_api_region()
    host = f"q.{region}.amazonaws.com"
    machine_id = generate_machine_id(effective_cfg)
    kiro_version = (effective_cfg.kiro_version or "0.8.0").strip() or "0.8.0"

    # 构建 URL（添加 isEmailRequired=true 获取邮箱）
    url = f"https://{host}/getUsageLimits?origin=AI_EDITOR&resourceType=AGENTIC_REQUEST&isEmailRequired=true"

    profile_arn = get_profile_arn_for_payload(effective_cfg)
    if profile_arn:
        from urllib.parse import quote

        url += f"&profileArn={quote(profile_arn, safe='')}"

    logger.debug("[KIRO_QUOTA] 请求 URL: {}", url)

    # 构建 headers
    headers = {
        "x-amz-user-agent": build_x_amz_user_agent_usage(
            kiro_version=kiro_version, machine_id=machine_id
        ),
        "User-Agent": build_user_agent_usage(kiro_version=kiro_version, machine_id=machine_id),
        "host": host,
        "amz-sdk-invocation-id": str(uuid.uuid4()),
        "amz-sdk-request": "attempt=1; max=1",
        "Authorization": f"Bearer {access_token}",
        "Connection": "close",
    }

    response = await execute_kiro_rust_http_request(
        method="GET",
        url=url,
        headers=headers,
        body=None,
        proxy_config=proxy_config,
        request_id=f"kiro-usage:{region}:{machine_id}",
        provider_api_format="kiro:usage",
    )
    if response is None:
        client = await HTTPClientPool.get_proxy_client(proxy_config=proxy_config)
        response = await client.get(url, headers=headers, timeout=httpx.Timeout(30.0))

    if response.status_code != 200:
        response_text = (response.text or "").strip()

        # 检测账户封禁（403/423 均视为封禁）
        # 对于 Kiro，403 表示账户权限被拒绝（无论是封禁、权限不足还是其他原因），
        # 都应该标记为异常状态，以便管理员及时处理
        is_banned = False
        ban_reason = None
        if response.status_code in (403, 423):
            is_banned = True
            # 检测封禁相关错误（用于提供更详细的原因说明）
            banned_keywords = [
                "AccountSuspendedException",
                "account.*suspend",
                "account.*banned",
                "account.*disabled",
                "account.*access.*denied",
            ]
            for keyword in banned_keywords:
                if re.search(keyword, response_text, re.IGNORECASE):
                    ban_reason = (
                        response_text[:200] if response_text else f"HTTP {response.status_code}"
                    )
                    break
            # 如果没有匹配到特定关键词，使用通用原因
            if not ban_reason:
                if response.status_code == 423:
                    ban_reason = response_text[:200] if response_text else "HTTP 423 Locked"
                else:
                    ban_reason = response_text[:200] if response_text else "HTTP 403 权限被拒绝"

        error_msg = {
            401: "认证失败，Token 无效或已过期",
            403: "账户异常，权限被拒绝",
            423: "账户已封禁",
            429: "请求过于频繁，已被限流",
        }.get(response.status_code, "获取使用额度失败")
        if 500 <= response.status_code < 600:
            error_msg = "服务器错误，AWS 服务暂时不可用"
        logger.debug(
            "kiro usage API error: HTTP {} | {}",
            response.status_code,
            response_text[:200],
        )

        # 如果检测到封禁，抛出带有封禁标记的异常
        if is_banned:
            raise KiroAccountBannedException(
                message=error_msg,
                status_code=response.status_code,
                reason=ban_reason,
            )

        raise RuntimeError(f"{error_msg}: HTTP {response.status_code}")

    try:
        data = response.json()
    except Exception as exc:
        raise RuntimeError(f"获取使用额度成功但响应解析失败: HTTP {response.status_code}") from exc

    # 返回刷新后的配置（用于更新 auth_config）
    return {
        "usage_data": data,
        "updated_auth_config": updated_cfg.to_dict() if updated_cfg else None,
    }


def parse_kiro_usage_response(data: dict) -> dict | None:
    """
    解析 Kiro getUsageLimits API 响应，提取限额信息和用户邮箱

    返回格式与 kiro.rs BalanceResponse 类似：
    - subscription_title: 订阅类型（如 "KIRO PRO+"）
    - current_usage: 当前使用量
    - usage_limit: 使用限额
    - remaining: 剩余额度
    - usage_percentage: 使用百分比
    - next_reset_at: 下次重置时间（Unix 时间戳）
    - email: 用户邮箱（通过 isEmailRequired=true 获取）
    """
    if not data:
        return None

    usage_resp = UsageLimitsResponse.from_dict(data)

    current_usage = calculate_current_usage(usage_resp)
    usage_limit = calculate_total_usage_limit(usage_resp)
    remaining = max(usage_limit - current_usage, 0.0)
    usage_percentage = (current_usage / usage_limit * 100.0) if usage_limit > 0 else 0.0
    usage_percentage = min(usage_percentage, 100.0)

    result: dict[str, Any] = {
        "current_usage": current_usage,
        "usage_limit": usage_limit,
        "remaining": remaining,
        "usage_percentage": usage_percentage,
    }

    # 订阅类型
    if usage_resp.subscription_info and usage_resp.subscription_info.subscription_title:
        result["subscription_title"] = usage_resp.subscription_info.subscription_title

    # 下次重置时间
    if usage_resp.next_date_reset is not None:
        result["next_reset_at"] = usage_resp.next_date_reset
    elif usage_resp.usage_breakdown_list and usage_resp.usage_breakdown_list[0].next_date_reset:
        result["next_reset_at"] = usage_resp.usage_breakdown_list[0].next_date_reset

    # 解析用户邮箱（从 desktopUserInfo 或 userInfo 中获取）
    user_info = data.get("desktopUserInfo") or data.get("userInfo") or {}
    if isinstance(user_info, dict):
        email = user_info.get("email")
        if isinstance(email, str) and email.strip():
            result["email"] = email.strip()

    # 添加更新时间戳
    result["updated_at"] = int(time.time())

    return result


__all__ = [
    "KiroAccountBannedException",
    "fetch_kiro_usage_limits",
    "parse_kiro_usage_response",
]