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

from __future__ import annotations

import asyncio
import hashlib
import hmac
import json
import time

import aiohttp

from yuxi.channels.exceptions import ChannelAuthenticationError
from yuxi.utils.logging_config import logger


class YuanbaoTokenManager:
    REFRESH_WINDOW_SECONDS = 300
    _local_cache: dict[str, tuple[str, float]] = {}
    _local_cache_lock = asyncio.Lock()

    def __init__(
        self,
        app_key: str,
        app_secret: str,
        bot_app_id: str,
        pre_signed_token: str | None = None,
        api_base: str | None = None,
        http_client: aiohttp.ClientSession | None = None,
    ):
        self._app_key = app_key
        self._app_secret = app_secret
        self._bot_app_id = bot_app_id
        self._pre_signed_token = pre_signed_token
        self._api_base = api_base or "https://open-api.yuanbao.tencent.com"
        self._access_token: str | None = pre_signed_token
        self._expires_at: float | None = None
        self._token_lock = asyncio.Lock()
        self._http_client = http_client

    @property
    def api_base(self) -> str:
        return self._api_base

    @property
    def bot_app_id(self) -> str:
        return self._bot_app_id

    def _sign_request(self, timestamp: int) -> str:
        message = f"{self._app_key}{timestamp}"
        return hmac.new(
            self._app_secret.encode("utf-8"),
            message.encode("utf-8"),
            hashlib.sha256,
        ).hexdigest()

    async def get_token(self) -> str:
        if self._pre_signed_token:
            return self._pre_signed_token
        async with self._token_lock:
            if self._is_expired():
                await self._refresh()
            return self._access_token

    def is_expired(self) -> bool:
        if self._access_token is None or self._expires_at is None:
            return True
        return time.time() > self._expires_at - self.REFRESH_WINDOW_SECONDS

    _is_expired = is_expired

    async def refresh_token(self) -> str:
        redis_key = f"yuanbao:token:{self._bot_app_id}"
        try:
            from yuxi.storage.redis import get_redis

            redis = get_redis()
            await redis.delete(redis_key)
        except Exception:
            logger.warning(f"[Yuanbao] Redis unavailable, skipping cache delete for {redis_key}")
        return await self.get_token()

    async def _refresh(self) -> None:
        if self._pre_signed_token:
            return
        redis_key = f"yuanbao:token:{self._bot_app_id}"
        cached_from_redis = False

        try:
            from yuxi.storage.redis import get_redis

            redis = get_redis()
            cached = await redis.get(redis_key)
            if cached:
                token_data = json.loads(cached)
                expires_at = token_data["expires_at"]
                if time.time() < expires_at - self.REFRESH_WINDOW_SECONDS:
                    self._access_token = token_data["access_token"]
                    self._expires_at = expires_at
                    cached_from_redis = True
                    return
        except Exception:
            logger.warning(f"[Yuanbao] Redis unavailable for read, checking local cache ({redis_key})")

        if not cached_from_redis:
            local_token = self._load_from_local_cache()
            if local_token is not None:
                return

        timestamp = int(time.time())
        signature = self._sign_request(timestamp)

        if self._http_client:
            session = self._http_client
            async with session.post(
                f"{self._api_base}/api/auth/token",
                json={
                    "app_key": self._app_key,
                    "timestamp": timestamp,
                    "signature": signature,
                },
                timeout=aiohttp.ClientTimeout(total=10),
            ) as resp:
                if resp.status != 200:
                    raise ChannelAuthenticationError(f"Token refresh failed: HTTP {resp.status} {await resp.text()}")
                data = await resp.json()
                self._access_token = data["access_token"]
                expires_in = data.get("expires_in", 7200)
                self._expires_at = time.time() + expires_in
        else:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    f"{self._api_base}/api/auth/token",
                    json={
                        "app_key": self._app_key,
                        "timestamp": timestamp,
                        "signature": signature,
                    },
                    timeout=aiohttp.ClientTimeout(total=10),
                ) as resp:
                    if resp.status != 200:
                        raise ChannelAuthenticationError(
                            f"Token refresh failed: HTTP {resp.status} {await resp.text()}"
                        )
                    data = await resp.json()
                    self._access_token = data["access_token"]
                    expires_in = data.get("expires_in", 7200)
                    self._expires_at = time.time() + expires_in

        self._save_to_local_cache(self._access_token)

        try:
            from yuxi.storage.redis import get_redis

            redis = get_redis()
            await redis.set(
                redis_key,
                json.dumps(
                    {
                        "access_token": self._access_token,
                        "expires_at": self._expires_at,
                        "obtained_at": time.time(),
                    }
                ),
                ex=expires_in,
            )
        except Exception:
            logger.warning(f"[Yuanbao] Redis unavailable for write, relying on local cache ({redis_key})")

        logger.info(f"[Yuanbao] Token refreshed, expires in {expires_in}s (bot_app_id={self._bot_app_id})")

    async def _load_from_local_cache(self) -> bool:
        async with self._local_cache_lock:
            entry = self._local_cache.get(self._bot_app_id)
            if entry is None:
                return False
            access_token, expires_at = entry
            if time.time() < expires_at - self.REFRESH_WINDOW_SECONDS:
                self._access_token = access_token
                self._expires_at = expires_at
                logger.info(
                    f"[Yuanbao] Token loaded from local cache (Redis unavailable) (bot_app_id={self._bot_app_id})"
                )
                return True
        return False

    def _save_to_local_cache(self, token: str) -> None:
        self._local_cache[self._bot_app_id] = (token, self._expires_at)