ForcePilot/backend/package/yuxi/channel/extensions/jd/gateway.py

import asyncio
import json
import logging
import os
import time

import httpx

from yuxi.channel.extensions.jd.config import JDConfig
from yuxi.channel.extensions.jd.types import JDAccount, JOSToken

logger = logging.getLogger(__name__)

OAUTH_TOKEN_URL = "https://oauth.jd.com/oauth/token"
TOKEN_REFRESH_MARGIN = 300
REFRESH_INTERVAL = 1800
RETRY_MAX = 3
RETRY_BASE_DELAY = 2

_TOKEN_STORE_DIR = os.path.join(os.path.dirname(__file__), ".token_cache")


class JDGateway:
    def __init__(self):
        self._account: JDAccount | None = None
        self._queue: asyncio.Queue | None = None
        self._cancel_event = asyncio.Event()
        self._access_token: str | None = None
        self._token_obj: JOSToken | None = None
        self._token_expires_at: float = 0
        self._refresh_task: asyncio.Task | None = None
        self._http: httpx.AsyncClient | None = None
        self._running = False
        self._token_lock = asyncio.Lock()
        self._last_chat_id: dict[str, str] = {}

    async def start(self, ctx) -> object:
        config = JDConfig()
        account = config.resolve_account(ctx.account_id if hasattr(ctx, "account_id") else "default")

        if not account.is_configured():
            logger.warning("JD account not configured, skipping start")
            return {"running": False, "reason": "not-configured"}

        self._account = account
        self._queue = asyncio.Queue(maxsize=1000)
        self._cancel_event.clear()
        self._http = httpx.AsyncClient(timeout=15.0)

        try:
            await self._restore_or_fetch_token()
            self._refresh_task = asyncio.create_task(self._token_refresh_loop())
            self._running = True
            logger.info(
                "JD gateway started: account=%s shop=%s token_expires=%d",
                account.account_id,
                account.shop_id,
                int(self._token_expires_at - time.time()) if self._token_expires_at else 0,
            )
            return {
                "running": True,
                "account_id": account.account_id,
                "shop_id": account.shop_id,
                "queue": self._queue,
            }
        except Exception:
            logger.exception("JD gateway failed to start")
            await self.stop(ctx)
            raise

    async def stop(self, ctx) -> None:
        self._running = False
        self._cancel_event.set()

        if self._refresh_task and not self._refresh_task.done():
            self._refresh_task.cancel()
            try:
                await self._refresh_task
            except asyncio.CancelledError:
                pass

        if self._http:
            await self._http.aclose()
            self._http = None

        self._queue = None
        self._last_chat_id.clear()
        logger.info("JD gateway stopped")

    @property
    def account(self) -> JDAccount | None:
        return self._account

    @property
    def queue(self) -> asyncio.Queue | None:
        return self._queue

    def get_token(self, account_id: str = "default") -> str | None:
        if self._access_token and time.time() < self._token_expires_at:
            return self._access_token
        return None

    def get_account(self, account_id: str = "default") -> JDAccount | None:
        return self._account

    def cache_chat_id(self, user_id: str, chat_id: str):
        self._last_chat_id[user_id] = chat_id

    def get_cached_chat_id(self, user_id: str) -> str | None:
        return self._last_chat_id.get(user_id)

    async def _restore_or_fetch_token(self):
        account = self._account

        loaded = await self._load_token_from_disk()
        if loaded and loaded.get("expires_at", 0) > time.time() + TOKEN_REFRESH_MARGIN:
            self._access_token = loaded["access_token"]
            self._token_expires_at = loaded["expires_at"]
            if loaded.get("refresh_token"):
                self._account.refresh_token = loaded["refresh_token"]
            logger.info("JD token restored from disk")
            return

        if account.access_token and account.token_expires_at:
            if account.token_expires_at > time.time() + TOKEN_REFRESH_MARGIN:
                self._access_token = account.access_token
                self._token_expires_at = account.token_expires_at
                return

        if account.refresh_token:
            try:
                await self._refresh_token()
            except Exception as e:
                logger.error("JD token refresh failed on start: %s", e)

        if not self._access_token:
            logger.warning("JD account %s has no valid access token", account.account_id)

    async def _token_refresh_loop(self):
        while not self._cancel_event.is_set():
            try:
                await asyncio.sleep(REFRESH_INTERVAL)
                if self._cancel_event.is_set():
                    return

                remaining = self._token_expires_at - time.time()
                if remaining < TOKEN_REFRESH_MARGIN:
                    logger.info("JD token refresh triggered: remaining=%ds", int(remaining))
                    await self._refresh_token()

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error("JD token refresh loop error: %s", e)

    async def _refresh_token(self):
        async with self._token_lock:
            if not self._account.refresh_token:
                return

            params = {
                "grant_type": "refresh_token",
                "client_id": self._account.app_key,
                "client_secret": self._account.app_secret,
                "refresh_token": self._account.refresh_token,
            }

            for attempt in range(RETRY_MAX):
                try:
                    resp = await self._http.post(OAUTH_TOKEN_URL, data=params)
                    resp.raise_for_status()
                    data = resp.json()

                    expires_in = data.get("expires_in", 7200)
                    self._access_token = data["access_token"]
                    self._token_expires_at = time.time() + expires_in - TOKEN_REFRESH_MARGIN

                    new_refresh = data.get("refresh_token", "")
                    if new_refresh:
                        self._account.refresh_token = new_refresh

                    self._token_obj = JOSToken(
                        access_token=data["access_token"],
                        refresh_token=data.get("refresh_token", self._account.refresh_token or ""),
                        expires_at=time.time() + expires_in,
                        uid=data.get("uid", ""),
                        user_nick=data.get("user_nick", ""),
                    )

                    await self._save_token_to_disk()

                    logger.info(
                        "JD token refreshed: expires_in=%ds uid=%s",
                        expires_in,
                        data.get("uid", ""),
                    )
                    return

                except Exception as e:
                    logger.warning("JD token refresh attempt %d failed: %s", attempt + 1, e)
                    if attempt == RETRY_MAX - 1:
                        raise
                    delay = RETRY_BASE_DELAY * (2**attempt)
                    await asyncio.sleep(delay)

    async def _save_token_to_disk(self):
        if not self._token_obj:
            return
        try:
            os.makedirs(_TOKEN_STORE_DIR, exist_ok=True)
            token_file = os.path.join(_TOKEN_STORE_DIR, f"{self._account.shop_id}.json")
            with open(token_file, "w") as f:
                json.dump(
                    {
                        "access_token": self._access_token,
                        "refresh_token": self._account.refresh_token,
                        "expires_at": self._token_expires_at,
                        "uid": self._token_obj.uid,
                        "user_nick": self._token_obj.user_nick,
                    },
                    f,
                )
        except Exception as e:
            logger.warning("Failed to save JD token to disk: %s", e)

    async def _load_token_from_disk(self) -> dict | None:
        if not self._account:
            return None
        token_file = os.path.join(_TOKEN_STORE_DIR, f"{self._account.shop_id}.json")
        if not os.path.exists(token_file):
            return None
        try:
            with open(token_file) as f:
                data = json.load(f)
            if data.get("expires_at", 0) > time.time():
                return data
        except Exception as e:
            logger.warning("Failed to load JD token from disk: %s", e)
        return None