Aether/_deprecated_py_src/services/provider/adapters/antigravity/client.py

"""Antigravity API 客户端（与 Antigravity-Manager 对齐）。"""

from __future__ import annotations

import asyncio
import json
import random
import re
from typing import Any

from src.clients.http_client import HTTPClientPool
from src.core.logger import logger
from src.services.provider.adapters.antigravity.constants import (
    DAILY_BASE_URL,
    PROD_BASE_URL,
    SANDBOX_BASE_URL,
    VERSION_FETCH_URL,
    get_v1internal_extra_headers,
    parse_version_string,
    update_user_agent_version,
)
from src.services.provider.adapters.antigravity.rust_http import (
    execute_antigravity_rust_http_request,
)
from src.services.provider.adapters.antigravity.url_availability import url_availability

# loadCodeAssist 请求体 metadata
_CODE_ASSIST_METADATA = {
    "ideType": "ANTIGRAVITY",
}

# Duration 解析正则（与 AM 的 retry.rs 对齐）
_DURATION_RE = re.compile(r"([\d.]+)\s*(ms|s|m|h)")


# ---------------------------------------------------------------------------
# Retry-After / Duration 解析工具（对齐 AM upstream/retry.rs）
# ---------------------------------------------------------------------------


def _parse_duration_ms(duration_str: str) -> int | None:
    """解析 Duration 字符串 (e.g. '1.5s', '200ms', '1h16m0.667s')，返回毫秒。"""
    total_ms = 0.0
    matched = False
    for m in _DURATION_RE.finditer(duration_str):
        matched = True
        value = float(m.group(1))
        unit = m.group(2)
        if unit == "ms":
            total_ms += value
        elif unit == "s":
            total_ms += value * 1000
        elif unit == "m":
            total_ms += value * 60_000
        elif unit == "h":
            total_ms += value * 3_600_000
    return round(total_ms) if matched else None


def parse_retry_delay(error_body: str | bytes | dict[str, Any]) -> float | None:
    """从 429 错误响应体中提取 retry delay（秒）。

    支持两种格式（与 AM 对齐）：
    1. error.details[].@type="...RetryInfo" → retryDelay
    2. error.details[].metadata.quotaResetDelay
    """
    try:
        data: dict[str, Any]
        if isinstance(error_body, (str, bytes)):
            data = json.loads(error_body)
        elif isinstance(error_body, dict):
            data = error_body
        else:
            return None

        details = data.get("error", {}).get("details", [])
        if not isinstance(details, list):
            return None

        # 方式 1: RetryInfo.retryDelay
        for detail in details:
            if not isinstance(detail, dict):
                continue
            type_str = detail.get("@type", "")
            if isinstance(type_str, str) and "RetryInfo" in type_str:
                retry_delay = detail.get("retryDelay")
                if isinstance(retry_delay, str):
                    ms = _parse_duration_ms(retry_delay)
                    if ms is not None:
                        return min((ms + 200) / 1000.0, 30.0)

        # 方式 2: metadata.quotaResetDelay
        for detail in details:
            if not isinstance(detail, dict):
                continue
            metadata = detail.get("metadata")
            if not isinstance(metadata, dict):
                continue
            quota_delay = metadata.get("quotaResetDelay")
            if isinstance(quota_delay, str):
                ms = _parse_duration_ms(quota_delay)
                if ms is not None:
                    return min((ms + 200) / 1000.0, 30.0)
    except Exception:
        pass
    return None


# ---------------------------------------------------------------------------
# Fallback 判断（对齐 AM upstream/client.rs: should_try_next_endpoint）
# ---------------------------------------------------------------------------


def _should_fallback_status(status_code: int) -> bool:
    """判断是否应该 fallback 到下一个端点。

    与 AM 对齐：429 + 408(超时) + 404 + 所有 5xx。
    4xx 客户端错误（400/401/403 等）不 fallback——换 URL 也不会成功。
    """
    if status_code == 429:
        return True
    if status_code in (404, 408):
        return True
    if 500 <= status_code < 600:  # noqa: PLR2004
        return True
    return False


# ---------------------------------------------------------------------------
# 账户封禁异常（对齐 AM: is_forbidden 标志）
# ---------------------------------------------------------------------------


class AntigravityAccountForbiddenException(Exception):
    """Antigravity 账户被封禁/禁止访问异常。

    当 API 返回 403 Forbidden 时抛出，表示账户权限被撤销。
    """

    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


def _extract_forbidden_reason(response_text: str) -> str | None:
    """从 403 响应体中提取封禁原因。

    尝试解析 JSON 响应中的 error.message 字段。
    """
    if not response_text:
        return None
    try:
        data = json.loads(response_text)
        if isinstance(data, dict):
            error = data.get("error")
            if isinstance(error, dict):
                message = error.get("message")
                if isinstance(message, str) and message.strip():
                    return message.strip()
            # 直接在顶层查找 message
            message = data.get("message")
            if isinstance(message, str) and message.strip():
                return message.strip()
    except Exception:
        pass
    # 如果无法解析，返回原始文本的前 100 个字符
    if len(response_text) > 100:
        return response_text[:100] + "..."
    return response_text if response_text.strip() else None


# ---------------------------------------------------------------------------
# 从 loadCodeAssist 响应中提取信息
# ---------------------------------------------------------------------------


def extract_tier_id(data: dict[str, Any]) -> str:
    """从 loadCodeAssist 响应中提取 tier ID。

    优先选 allowedTiers 中 isDefault=true 的，fallback 到第一个，最终 fallback 到 "LEGACY"。
    """
    allowed_tiers = data.get("allowedTiers")
    if not isinstance(allowed_tiers, list):
        return "LEGACY"

    # 第一轮：找 isDefault
    for tier in allowed_tiers:
        if isinstance(tier, dict) and tier.get("isDefault") is True:
            tier_id = tier.get("id", "")
            if isinstance(tier_id, str) and tier_id.strip():
                return tier_id.strip()

    # 第二轮：取第一个有 id 的
    for tier in allowed_tiers:
        if isinstance(tier, dict):
            tier_id = tier.get("id", "")
            if isinstance(tier_id, str) and tier_id.strip():
                return tier_id.strip()

    return "LEGACY"


def extract_project_id(data: dict[str, Any]) -> str:
    """从响应中提取 project_id，兼容 string 和 {"id": "..."} 两种格式。"""
    raw = data.get("cloudaicompanionProject")
    if isinstance(raw, str) and raw.strip():
        return raw.strip()
    if isinstance(raw, dict):
        pid = raw.get("id", "")
        if isinstance(pid, str) and pid.strip():
            return pid.strip()
    return ""


# ---------------------------------------------------------------------------
# 核心 API 客户端函数
# ---------------------------------------------------------------------------


async def load_code_assist(
    access_token: str,
    proxy_config: dict[str, Any] | None = None,
    *,
    timeout_seconds: float = 10.0,
) -> dict[str, Any]:
    """调用 /v1internal:loadCodeAssist 获取账户信息。

    对齐 AM project_resolver.rs：Sandbox 优先，避免 Prod 429。

    注意：
    - email 需通过 Google userinfo API 获取（由 enrich_auth_config 复用已有逻辑）
    - 这里仅负责 project_id / tier 等信息
    """
    if not access_token:
        raise ValueError("missing access_token")

    headers = {
        "Authorization": f"Bearer {access_token}",
        "Content-Type": "application/json",
        **get_v1internal_extra_headers(),
    }
    body = {"metadata": _CODE_ASSIST_METADATA}

    # Sandbox 优先（与 AM 对齐：避免 Prod 429）
    urls = url_availability.get_ordered_urls(prefer_daily=True)
    if not urls:
        urls = [SANDBOX_BASE_URL, DAILY_BASE_URL, PROD_BASE_URL]
    last_exc: Exception | None = None

    for base_url in urls:
        try:
            url = f"{base_url}/v1internal:loadCodeAssist"
            resp = await execute_antigravity_rust_http_request(
                method="POST",
                url=url,
                headers=headers,
                body=body,
                proxy_config=proxy_config,
                request_id=f"antigravity:load-code-assist:{base_url}",
                provider_api_format="antigravity:load_code_assist",
                timeout_seconds=timeout_seconds,
                content_type="application/json",
            )
            if resp is None:
                client = await HTTPClientPool.get_proxy_client(proxy_config)
                resp = await client.post(
                    url,
                    json=body,
                    headers=headers,
                    timeout=timeout_seconds,
                )
            if 200 <= resp.status_code < 300:
                url_availability.mark_success(base_url)
                data = resp.json()
                return data if isinstance(data, dict) else {}

            # 可 fallback 的错误：标记不可用，尝试下一个 URL
            if _should_fallback_status(resp.status_code):
                url_availability.mark_unavailable(base_url)

                # 429：尝试解析 Retry-After 并等待
                if resp.status_code == 429:
                    delay = parse_retry_delay(resp.text)
                    if delay and delay > 0:
                        logger.debug(
                            "[antigravity] loadCodeAssist 429, waiting {:.1f}s before fallback",
                            delay,
                        )
                        await asyncio.sleep(min(delay, 5.0))

                last_exc = RuntimeError(
                    f"loadCodeAssist failed: status={resp.status_code} base_url={base_url}"
                )
                continue

            # 不可 fallback 的 4xx（400/401/403 等）：直接抛出，换 URL 也不会成功
            raise RuntimeError(
                f"loadCodeAssist failed: status={resp.status_code} base_url={base_url} "
                f"body={resp.text[:200] if resp.text else ''}"
            )
        except RuntimeError:
            raise
        except Exception as e:
            url_availability.mark_unavailable(base_url)
            last_exc = e
            continue

    raise last_exc or RuntimeError("loadCodeAssist failed: all endpoints exhausted")


async def onboard_user(
    access_token: str,
    tier_id: str = "LEGACY",
    proxy_config: dict[str, Any] | None = None,
    *,
    max_attempts: int = 5,
    poll_interval: float = 2.0,
    timeout_seconds: float = 30.0,
) -> str:
    """调用 /v1internal:onboardUser 激活账号并获取 project_id。

    当 loadCodeAssist 返回 allowedTiers 但没有 cloudaicompanionProject 时，
    需要先通过 onboardUser 选择 tier 来分配 project。

    改进（对齐 AM）：
    - 使用 url_availability 做多 URL fallback
    - 轮询期间的临时网络错误会 continue 而非直接终止

    Args:
        access_token: OAuth access token
        tier_id: 要选择的 tier ID（从 allowedTiers 中提取）
        proxy_config: 代理配置
        max_attempts: 最大轮询次数（onboardUser 是异步操作）
        poll_interval: 轮询间隔（秒）
        timeout_seconds: 单次请求超时

    Returns:
        project_id

    Raises:
        RuntimeError: 激活失败
    """
    if not access_token:
        raise ValueError("missing access_token")

    headers = {
        "Authorization": f"Bearer {access_token}",
        "Content-Type": "application/json",
        **get_v1internal_extra_headers(),
    }
    body = {
        "tierId": tier_id,
        "metadata": _CODE_ASSIST_METADATA,
    }

    # 使用 url_availability 选择端点（与其他函数一致）
    urls = url_availability.get_ordered_urls(prefer_daily=True)
    if not urls:
        urls = [SANDBOX_BASE_URL, DAILY_BASE_URL, PROD_BASE_URL]
    base_url = urls[0]

    logger.info("[antigravity] onboardUser: 开始激活账号, tier={}, endpoint={}", tier_id, base_url)

    for attempt in range(1, max_attempts + 1):
        try:
            url = f"{base_url}/v1internal:onboardUser"
            resp = await execute_antigravity_rust_http_request(
                method="POST",
                url=url,
                headers=headers,
                body=body,
                proxy_config=proxy_config,
                request_id=f"antigravity:onboard-user:{tier_id}:{attempt}",
                provider_api_format="antigravity:onboard_user",
                timeout_seconds=timeout_seconds,
                content_type="application/json",
            )
            if resp is None:
                client = await HTTPClientPool.get_proxy_client(proxy_config)
                resp = await client.post(
                    url,
                    json=body,
                    headers=headers,
                    timeout=timeout_seconds,
                )
            if resp.status_code < 200 or resp.status_code >= 300:
                # 标记可用性
                if _should_fallback_status(resp.status_code):
                    url_availability.mark_unavailable(base_url)
                raise RuntimeError(
                    f"onboardUser failed: status={resp.status_code}, "
                    f"body={resp.text[:200] if resp.text else ''}"
                )

            url_availability.mark_success(base_url)
            data = resp.json()
            if not isinstance(data, dict):
                raise RuntimeError(f"onboardUser: unexpected response type: {type(data)}")

            done = data.get("done")
            if done is True:
                # 从 response.cloudaicompanionProject 提取 project_id
                response_data = data.get("response")
                if isinstance(response_data, dict):
                    project_id = extract_project_id(response_data)
                    if project_id:
                        logger.info(
                            "[antigravity] onboardUser: 激活成功, project_id={}",
                            project_id[:8] + "...",
                        )
                        return project_id

                # done=true 但 project_id 为空（常见：cloudaicompanionProject: {}）
                # 返回空串，由调用方 fallback 到随机 project_id
                logger.debug("[antigravity] onboardUser: done=true 但 project_id 为空")
                return ""

            # done != true，继续轮询
            logger.debug(
                "[antigravity] onboardUser: 轮询 {}/{}, 等待完成...",
                attempt,
                max_attempts,
            )
            if attempt < max_attempts:
                await asyncio.sleep(poll_interval)

        except RuntimeError:
            raise
        except Exception as e:
            # 临时网络错误：记录并继续轮询（而非直接终止）
            logger.warning(
                "[antigravity] onboardUser: 轮询 {}/{} 网络错误: {}, 继续重试...",
                attempt,
                max_attempts,
                e,
            )
            if attempt < max_attempts:
                await asyncio.sleep(poll_interval)
                continue
            raise RuntimeError(
                f"onboardUser request failed after {max_attempts} attempts: {e}"
            ) from e

    raise RuntimeError(f"onboardUser: 超时，已轮询 {max_attempts} 次仍未完成")


async def fetch_available_models(
    access_token: str,
    *,
    project_id: str,
    proxy_config: dict[str, Any] | None = None,
    timeout_seconds: float = 10.0,
) -> dict[str, Any]:
    """调用 /v1internal:fetchAvailableModels 获取可用模型（包含配额信息）。

    Antigravity 的此接口会返回类似：
        {"models": {"claude-sonnet-4": {"displayName": "...", "quotaInfo": {...}}, ...}}

    对齐 AM：Sandbox 优先 + 正确的 fallback 逻辑。
    """
    if not access_token:
        raise ValueError("missing access_token")
    if not project_id:
        raise ValueError("missing project_id")

    headers = {
        "Authorization": f"Bearer {access_token}",
        "Content-Type": "application/json",
        "Accept": "application/json",
        **get_v1internal_extra_headers(),
    }
    body = {"project": project_id}

    # Sandbox 优先
    urls = url_availability.get_ordered_urls(prefer_daily=True)
    if not urls:
        urls = [SANDBOX_BASE_URL, DAILY_BASE_URL, PROD_BASE_URL]
    last_exc: Exception | None = None

    for base_url in urls:
        try:
            url = f"{base_url}/v1internal:fetchAvailableModels"
            resp = await execute_antigravity_rust_http_request(
                method="POST",
                url=url,
                headers=headers,
                body=body,
                proxy_config=proxy_config,
                request_id=f"antigravity:fetch-available-models:{base_url}",
                provider_api_format="antigravity:fetch_available_models",
                timeout_seconds=timeout_seconds,
                content_type="application/json",
            )
            if resp is None:
                client = await HTTPClientPool.get_proxy_client(proxy_config)
                resp = await client.post(
                    url,
                    json=body,
                    headers=headers,
                    timeout=timeout_seconds,
                )
            if 200 <= resp.status_code < 300:
                url_availability.mark_success(base_url)
                data = resp.json()
                return data if isinstance(data, dict) else {}

            # 可 fallback 的错误
            if _should_fallback_status(resp.status_code):
                url_availability.mark_unavailable(base_url)

                # 429：尝试等待
                if resp.status_code == 429:
                    delay = parse_retry_delay(resp.text)
                    if delay and delay > 0:
                        logger.debug(
                            "[antigravity] fetchAvailableModels 429, waiting {:.1f}s",
                            delay,
                        )
                        await asyncio.sleep(min(delay, 5.0))

                last_exc = RuntimeError(
                    f"fetchAvailableModels failed: status={resp.status_code} base_url={base_url}"
                )
                continue

            # 403 Forbidden：账户权限被禁止（对齐 AM is_forbidden 标志）
            if resp.status_code == 403:
                reason = _extract_forbidden_reason(resp.text)
                logger.warning(
                    "[antigravity] fetchAvailableModels 403 Forbidden: {}",
                    reason or "unknown",
                )
                raise AntigravityAccountForbiddenException(
                    message="账户访问被禁止",
                    status_code=403,
                    reason=reason,
                )

            # 不可 fallback 的 4xx：直接抛出
            raise RuntimeError(
                f"fetchAvailableModels failed: status={resp.status_code} base_url={base_url} "
                f"body={resp.text[:200] if resp.text else ''}"
            )
        except (RuntimeError, AntigravityAccountForbiddenException):
            raise
        except Exception as e:
            url_availability.mark_unavailable(base_url)
            last_exc = e
            continue

    raise last_exc or RuntimeError("fetchAvailableModels failed: all endpoints exhausted")


# ---------------------------------------------------------------------------
# User-Agent 版本号动态更新
# ---------------------------------------------------------------------------


async def refresh_user_agent(
    proxy_config: dict[str, Any] | None = None,
    *,
    timeout_seconds: float = 5.0,
) -> str | None:
    """从远程 API 获取最新 Antigravity 版本号并更新 User-Agent。

    对齐 AM constants.rs：
    1. 尝试 VERSION_FETCH_URL
    2. Fallback 到 changelog 页面
    3. 最终 fallback 到 _FALLBACK_VERSION

    Returns:
        更新后的版本号，或 None（如果获取失败但 fallback 到默认版本）。
    """
    try:
        client = await HTTPClientPool.get_proxy_client(proxy_config)
        resp = await client.get(VERSION_FETCH_URL, timeout=timeout_seconds)
        if 200 <= resp.status_code < 300:
            version = parse_version_string(resp.text)
            if version:
                update_user_agent_version(version)
                logger.info("[antigravity] User-Agent 版本已更新: {}", version)
                return version
    except Exception as e:
        logger.debug("[antigravity] 获取远程版本失败: {}", e)

    return None


# ---------------------------------------------------------------------------
# Fallback Project ID 生成
# ---------------------------------------------------------------------------

_ADJECTIVES = ("useful", "bright", "swift", "calm", "bold")
_NOUNS = ("fuze", "wave", "spark", "flow", "core")


def generate_fallback_project_id() -> str:
    """生成随机 project_id（与 AM project_resolver.rs 的 generateProjectID 一致）。

    格式: {adjective}-{noun}-{5位随机字符(base36)}
    Antigravity API 不严格校验 project 字段，当所有正常获取途径都失败时用作 fallback。
    """
    adj = random.choice(_ADJECTIVES)  # noqa: S311
    noun = random.choice(_NOUNS)  # noqa: S311
    chars = "abcdefghijklmnopqrstuvwxyz0123456789"
    rand_part = "".join(random.choice(chars) for _ in range(5))  # noqa: S311
    return f"{adj}-{noun}-{rand_part}"


__all__ = [
    "AntigravityAccountForbiddenException",
    "extract_project_id",
    "extract_tier_id",
    "fetch_available_models",
    "generate_fallback_project_id",
    "load_code_assist",
    "onboard_user",
    "parse_retry_delay",
    "refresh_user_agent",
]