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

import asyncio
import logging
import time

import httpx

logger = logging.getLogger(__name__)

TOKEN_URL = "https://open.douyin.com/oauth/client_token/"
BUSINESS_TOKEN_URL = "https://open.douyin.com/oauth/business_token/"
REFRESH_BIZ_TOKEN_URL = "https://open.douyin.com/oauth/refresh_biz_token/"
REFRESH_INTERVAL = 3600
BIZ_REFRESH_INTERVAL = 29 * 86400
RETRY_MAX = 3
RETRY_BASE_DELAY = 2


class DouyinGateway:
    def __init__(self, config_adapter=None):
        self._config = config_adapter
        self._account = None
        self._client_token: str | None = None
        self._token_expires_at: float = 0
        self._refresh_task: asyncio.Task | None = None
        self._business_token: str | None = None
        self._biz_token_expires_at: float = 0
        self._biz_refresh_token: str | None = None
        self._biz_refresh_task: asyncio.Task | None = None
        self._cancel_event = asyncio.Event()
        self.http: httpx.AsyncClient | None = None
        self._running = False
        self._token_lock = asyncio.Lock()
        self._biz_token_lock = asyncio.Lock()

    async def start(self, ctx) -> object:
        from yuxi.channel.extensions.douyin.config import DouyinConfig

        config_adapter = self._config or DouyinConfig()
        account = config_adapter.resolve_account()
        if not account.is_configured():
            logger.warning("Douyin account not configured, skipping start")
            return {"running": False, "reason": "not-configured"}

        self._account = account
        self._cancel_event.clear()
        self.http = httpx.AsyncClient(
            base_url="https://open.douyin.com",
            timeout=15.0,
        )

        try:
            await self._fetch_client_token()
            await self._fetch_business_token()

            self._refresh_task = asyncio.create_task(self._token_refresh_loop())
            self._biz_refresh_task = asyncio.create_task(self._biz_token_refresh_loop())
            self._running = True
            logger.info("Douyin gateway started for account %s", account.account_id)
            return {"running": True, "account_id": account.account_id}
        except Exception:
            logger.exception("Douyin gateway failed to start")
            await self.stop(ctx)
            raise

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

        for task in (self._refresh_task, self._biz_refresh_task):
            if task and not task.done():
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass

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

        logger.info("Douyin gateway stopped")

    @property
    def client_token(self) -> str | None:
        if self._client_token and time.time() < self._token_expires_at:
            return self._client_token
        return None

    @property
    def business_token(self) -> str | None:
        if self._business_token and time.time() < self._biz_token_expires_at:
            return self._business_token
        return None

    async def _fetch_client_token(self) -> str:
        async with self._token_lock:
            if self._client_token and time.time() < self._token_expires_at:
                return self._client_token

            for attempt in range(RETRY_MAX):
                try:
                    resp = await self.http.post(
                        "/oauth/client_token/",
                        json={
                            "client_key": self._account.client_key,
                            "client_secret": self._account.client_secret,
                            "grant_type": "client_credential",
                        },
                    )
                    data = resp.json()
                    inner = data.get("data", {})
                    if inner.get("error_code", -1) != 0:
                        raise ConnectionError(f"Failed to get client_token: {data}")

                    self._client_token = inner["access_token"]
                    expires_in = inner.get("expires_in", 7200)
                    self._token_expires_at = time.time() + expires_in - 300
                    logger.info("Douyin client_token refreshed, expires_in=%ds", expires_in)
                    return self._client_token

                except Exception:
                    if attempt == RETRY_MAX - 1:
                        raise
                delay = RETRY_BASE_DELAY * (2**attempt)
                await asyncio.sleep(delay)

            raise ConnectionError("Failed to get client_token after retries")

    async def _fetch_business_token(self) -> str:
        async with self._biz_token_lock:
            if self._business_token and time.time() < self._biz_token_expires_at:
                return self._business_token

            client_token = self._client_token
            if not client_token:
                client_token = await self._fetch_client_token()

            for attempt in range(RETRY_MAX):
                try:
                    resp = await self.http.post(
                        "/oauth/business_token/",
                        json={
                            "client_key": self._account.client_key,
                            "client_secret": self._account.client_secret,
                            "grant_type": "client_credential",
                        },
                        headers={"access-token": client_token},
                    )
                    data = resp.json()
                    inner = data.get("data", {})
                    if inner.get("error_code", -1) != 0:
                        err_msg = data.get("message", inner.get("description", str(data)))
                        logger.error("Failed to get business_token: %s", err_msg)
                        if attempt == RETRY_MAX - 1:
                            raise ConnectionError(f"Failed to get business_token: {data}")
                        delay = RETRY_BASE_DELAY * (2**attempt)
                        await asyncio.sleep(delay)
                        continue

                    self._business_token = inner["access_token"]
                    self._biz_refresh_token = inner.get("refresh_token")
                    expires_in = inner.get("expires_in", 2592000)
                    self._biz_token_expires_at = time.time() + expires_in - 86400
                    logger.info(
                        "Douyin business_token refreshed, expires_in=%ds, has_refresh=%s",
                        expires_in,
                        bool(self._biz_refresh_token),
                    )
                    return self._business_token

                except Exception:
                    if attempt == RETRY_MAX - 1:
                        raise
                delay = RETRY_BASE_DELAY * (2**attempt)
                await asyncio.sleep(delay)

            raise ConnectionError("Failed to get business_token after retries")

    async def _refresh_biz_token(self) -> str:
        if not self._biz_refresh_token:
            return await self._fetch_business_token()

        for attempt in range(RETRY_MAX):
            try:
                resp = await self.http.post(
                    "/oauth/refresh_biz_token/",
                    json={
                        "client_key": self._account.client_key,
                        "grant_type": "refresh_biz_token",
                        "refresh_token": self._biz_refresh_token,
                    },
                )
                data = resp.json()
                inner = data.get("data", {})
                if inner.get("error_code", -1) != 0:
                    logger.warning("business_token refresh failed, retrying full fetch")
                    return await self._fetch_business_token()

                self._business_token = inner["access_token"]
                self._biz_refresh_token = inner.get("refresh_token")
                expires_in = inner.get("expires_in", 2592000)
                self._biz_token_expires_at = time.time() + expires_in - 86400
                logger.info("Douyin business_token refreshed via refresh_token, expires_in=%ds", expires_in)
                return self._business_token

            except Exception:
                if attempt == RETRY_MAX - 1:
                    logger.exception("business_token refresh exhausted retries")
                    return await self._fetch_business_token()
            delay = RETRY_BASE_DELAY * (2**attempt)
            await asyncio.sleep(delay)

        return await self._fetch_business_token()

    async def _token_refresh_loop(self):
        while not self._cancel_event.is_set():
            try:
                await asyncio.sleep(REFRESH_INTERVAL)
                if not self._cancel_event.is_set():
                    await self._fetch_client_token()
            except asyncio.CancelledError:
                break
            except Exception:
                logger.exception("Douyin client_token refresh failed in loop")

    async def _biz_token_refresh_loop(self):
        while not self._cancel_event.is_set():
            try:
                await asyncio.sleep(BIZ_REFRESH_INTERVAL)
                if not self._cancel_event.is_set():
                    async with self._biz_token_lock:
                        await self._refresh_biz_token()
            except asyncio.CancelledError:
                break
            except Exception:
                logger.exception("Douyin business_token refresh failed in loop")