ForcePilot/backend/package/yuxi/channel/extensions/yuanbao/accounts.py

from __future__ import annotations

import hashlib
import hmac
import logging
import os
import time
from dataclasses import dataclass

import httpx

from yuxi.channel.extensions.yuanbao.types import YuanbaoAccountConfig

logger = logging.getLogger(__name__)

_TOKEN_URL = "https://{api_domain}/openapi/v1/token"
_DEFAULT_API_DOMAIN = "bot.yuanbao.tencent.com"
_DEFAULT_WS_URL = "wss://bot-wss.yuanbao.tencent.com/wss/connection"


@dataclass
class ResolvedYuanbaoAccount:
    account_id: str
    app_key: str
    app_secret: str
    name: str | None = None
    api_domain: str = _DEFAULT_API_DOMAIN
    ws_url: str = _DEFAULT_WS_URL
    enabled: bool = True


async def resolve_yuanbao_account(account_id: str, config: dict) -> ResolvedYuanbaoAccount:
    account = config.get("channels", {}).get("yuanbao", {}).get("accounts", {}).get(account_id, {})

    app_key = account.get("appKey") or os.environ.get("YUANBAO_APP_KEY", "")
    app_secret = account.get("appSecret") or os.environ.get("YUANBAO_APP_SECRET", "")
    token = account.get("token")

    if (not app_key or not app_secret) and token:
        colon_idx = token.find(":")
        if colon_idx > 0 and colon_idx < len(token) - 1:
            parsed_key = token[:colon_idx].strip()
            parsed_secret = token[colon_idx + 1:].strip()
            if parsed_key and parsed_secret:
                app_key = app_key or parsed_key
                app_secret = app_secret or parsed_secret

    return ResolvedYuanbaoAccount(
        account_id=account_id,
        app_key=app_key,
        app_secret=app_secret,
        name=account.get("name"),
        api_domain=account.get("apiDomain", _DEFAULT_API_DOMAIN),
        ws_url=account.get("wsUrl", _DEFAULT_WS_URL),
        enabled=account.get("enabled", True),
    )


def _build_hmac_signature(app_key: str, app_secret: str) -> tuple[str, str]:
    timestamp = str(int(time.time()))
    signature_raw = f"{app_key}{timestamp}{app_secret}"
    signature = hmac.new(
        app_secret.encode("utf-8"),
        signature_raw.encode("utf-8"),
        hashlib.sha256,
    ).hexdigest()
    return timestamp, signature


async def fetch_access_token(
    app_key: str,
    app_secret: str,
    api_domain: str = _DEFAULT_API_DOMAIN,
) -> str:
    timestamp, signature = _build_hmac_signature(app_key, app_secret)
    url = _TOKEN_URL.format(api_domain=api_domain)

    async with httpx.AsyncClient() as client:
        resp = await client.post(
            url,
            json={"appKey": app_key, "signature": signature, "timestamp": timestamp},
            headers={"Content-Type": "application/json"},
            timeout=30.0,
        )
        resp.raise_for_status()
        data = resp.json()
        return data["access_token"]


class TokenCache:
    TTL_SECONDS = 2 * 60 * 60 - 60

    def __init__(self):
        self._tokens: dict[str, tuple[str, float]] = {}

    async def get(self, account: ResolvedYuanbaoAccount) -> str:
        key = account.account_id
        now = time.time()

        if key in self._tokens:
            token, expiry = self._tokens[key]
            if now < expiry:
                return token

        token = await fetch_access_token(
            account.app_key, account.app_secret, account.api_domain
        )
        self._tokens[key] = (token, now + self.TTL_SECONDS)
        return token

    def invalidate(self, account_id: str) -> None:
        self._tokens.pop(account_id, None)


def build_account_config(raw_account: dict, account_id: str) -> YuanbaoAccountConfig:
    return YuanbaoAccountConfig(
        account_id=account_id,
        app_key=raw_account.get("appKey", ""),
        app_secret=raw_account.get("appSecret", ""),
        name=raw_account.get("name"),
        api_domain=raw_account.get("apiDomain", _DEFAULT_API_DOMAIN),
        ws_url=raw_account.get("wsUrl", _DEFAULT_WS_URL),
        enabled=raw_account.get("enabled", True),
        max_chars=raw_account.get("maxChars", 3000),
        merge_text=raw_account.get("mergeText", "merge-text"),
        merge_text_min_chars=raw_account.get("mergeTextMinChars", 2800),
        merge_text_idle_ms=raw_account.get("mergeTextIdleMs", 5000),
        merge_text_overflow=raw_account.get("mergeTextOverflow", "split"),
        disable_block_streaming=raw_account.get("disableBlockStreaming", False),
        block_streaming_coalesce_min_chars=raw_account.get("blockStreamingCoalesceMinChars", 2800),
        block_streaming_coalesce_max_chars=raw_account.get("blockStreamingCoalesceMaxChars", 3000),
        block_streaming_coalesce_idle_ms=raw_account.get("blockStreamingCoalesceIdleMs", 1000),
        reply_to_mode=raw_account.get("replyToMode", "first"),
        dm_policy=raw_account.get("dmPolicy", "open"),
        pairing_ttl_seconds=raw_account.get("pairingTtlSeconds", 3600),
        session_key_ttl_seconds=raw_account.get("sessionKeyTtlSeconds", 300),
        history_limit=raw_account.get("historyLimit", 100),
        group_mention_required=raw_account.get("groupMentionRequired", True),
        debug_whitelist=raw_account.get("debugWhitelist", []),
        media_max_mb=raw_account.get("mediaMaxMb", 20),
        default_account=raw_account.get("defaultAccount", False),
        raw=raw_account,
    )