ForcePilot/backend/package/yuxi/channels/adapters/yuanbao/setup.py

from __future__ import annotations

from dataclasses import dataclass, field

from yuxi.utils.logging_config import logger


@dataclass
class SetupStepResult:
    step: str
    success: bool
    message: str = ""
    data: dict | None = None


@dataclass
class SetupWizardResult:
    app_key: str
    app_secret: str
    bot_app_id: str
    verified: bool = False
    error: str | None = None
    bot_info: dict | None = None
    steps: list[SetupStepResult] = field(default_factory=list)

    def add_step(self, step: str, success: bool, message: str = "", data: dict | None = None) -> None:
        self.steps.append(SetupStepResult(step=step, success=success, message=message, data=data))


async def run_setup_wizard(
    token_manager,
    http_client,
    app_key: str,
    app_secret: str,
    bot_app_id: str = "",
) -> SetupWizardResult:
    result = SetupWizardResult(
        app_key=app_key,
        app_secret=app_secret,
        bot_app_id=bot_app_id,
    )

    if not app_key.strip() or not app_secret.strip():
        result.error = "app_key 和 app_secret 不能为空"
        result.add_step("input_validation", False, result.error)
        return result

    result.add_step("input_validation", True, "输入参数校验通过")

    try:
        import time
        import hashlib
        import hmac
        import aiohttp

        timestamp = int(time.time())
        message = f"{app_key}{timestamp}"
        signature = hmac.new(
            app_secret.encode("utf-8"),
            message.encode("utf-8"),
            hashlib.sha256,
        ).hexdigest()

        api_base = "https://open-api.yuanbao.tencent.com"

        async with aiohttp.ClientSession() as session:
            async with session.post(
                f"{api_base}/api/auth/token",
                json={
                    "app_key": app_key,
                    "timestamp": timestamp,
                    "signature": signature,
                },
                timeout=aiohttp.ClientTimeout(total=10),
            ) as resp:
                if resp.status != 200:
                    result.error = f"认证失败: HTTP {resp.status}"
                    result.add_step("auth", False, result.error)
                    return result
                token_data = await resp.json()
                token = token_data.get("access_token", "")
                expires_in = token_data.get("expires_in", 7200)
                result.add_step("auth", True, f"认证成功，token 有效期 {expires_in}s")

            headers = {"Authorization": f"Bearer {token}"}
            async with session.get(
                f"{api_base}/api/v1/bot/info",
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=10),
            ) as resp:
                if resp.status == 200:
                    bot_info = await resp.json()
                    result.verified = True
                    result.bot_info = bot_info
                    result.bot_app_id = bot_info.get("bot_app_id", bot_app_id)
                    result.add_step(
                        "bot_verify",
                        True,
                        f"Bot 验证成功: {result.bot_app_id}",
                        bot_info,
                    )
                    logger.info(f"[Yuanbao] Setup wizard verified: bot_app_id={result.bot_app_id}")
                elif resp.status in (401, 403):
                    result.error = "凭据验证失败，请检查 app_key 和 app_secret"
                    result.add_step("bot_verify", False, result.error)
                else:
                    result.error = f"Bot 信息获取失败: HTTP {resp.status}"
                    result.add_step("bot_verify", False, result.error)

    except Exception as e:
        result.error = f"设置向导出错: {e}"
        result.add_step("unexpected_error", False, result.error)
        logger.error(f"[Yuanbao] Setup wizard error: {e}")

    return result


def generate_config_snippet(app_key: str, app_secret: str, bot_app_id: str) -> str:
    return f"""```yaml
channels:
  yuanbao:
    enabled: true
    app_key: "{app_key}"
    app_secret: "{app_secret}"
    bot_app_id: "{bot_app_id}"
    dm_policy: open
    requireMention: true
    overflowPolicy: split
    replyToMode: first
```"""


def generate_config_snippet_with_accounts(
    app_key: str,
    app_secret: str,
    bot_app_id: str,
    accounts: list[dict] | None = None,
) -> str:
    accounts_config = ""
    if accounts:
        for acct in accounts:
            accounts_config += f"""
    accounts.{acct["id"]}:
      appKey: "{acct.get("app_key", "")}"
      appSecret: "{acct.get("app_secret", "")}"
      name: "{acct.get("name", acct["id"])}"
      enabled: true"""

    return f"""```yaml
channels:
  yuanbao:
    enabled: true
    app_key: "{app_key}"
    app_secret: "{app_secret}"
    bot_app_id: "{bot_app_id}"
    dm_policy: open
    requireMention: true
    overflowPolicy: split
    replyToMode: first
    defaultAccount: "default"{accounts_config}
```"""