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

import asyncio
import logging
import time

import httpx

logger = logging.getLogger(__name__)

TOKEN_URL = "https://qyapi.weixin.qq.com/cgi-bin/gettoken"
REFRESH_INTERVAL = 7000
RETRY_MAX = 3
RETRY_BASE_DELAY = 2

_gateway: "WeComGateway | None" = None


def _set_gateway(gw: "WeComGateway") -> None:
    global _gateway
    _gateway = gw


def _get_gateway() -> "WeComGateway | None":
    return _gateway


class WeComGateway:
    def __init__(self):
        self._account = None
        self._cancel_event = asyncio.Event()
        self._access_token: str | 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()

    async def start(self, ctx) -> object:
        account = self._resolve_account(ctx)
        if not account.is_configured():
            logger.warning("WeCom account not configured, skipping start")
            return {"running": False, "reason": "not-configured"}

        self._account = account
        self._cancel_event.clear()
        self._http = httpx.AsyncClient(timeout=15.0)

        try:
            await self._refresh_token()
            self._refresh_task = asyncio.create_task(self._token_refresh_loop())
            self._running = True
            logger.info("WeCom gateway started for account %s", account.account_id)
            return {"running": True, "account_id": account.account_id}
        except Exception:
            logger.exception("WeCom 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

        logger.info("WeCom gateway stopped")

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

    @property
    def account(self):
        return self._account

    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._refresh_token()
            except asyncio.CancelledError:
                break

    async def _refresh_token(self):
        async with self._token_lock:
            if self._access_token and time.time() < self._token_expires_at:
                return

            url = f"{TOKEN_URL}?corpid={self._account.corp_id}&corpsecret={self._account.corp_secret}"

            for attempt in range(RETRY_MAX):
                try:
                    resp = await self._http.get(url)
                    data = resp.json()

                    if "access_token" in data:
                        self._access_token = data["access_token"]
                        self._token_expires_at = time.time() + data.get("expires_in", 7200) - 200
                        logger.info("WeCom access_token refreshed, expires in %ds", data["expires_in"])
                        return

                    errcode = data.get("errcode")
                    logger.error("Token refresh failed: errcode=%s, errmsg=%s", errcode, data.get("errmsg"))

                    if errcode in (-1, 40001, 40014):
                        raise ConnectionError(f"Invalid credentials: {data.get('errmsg')}")

                except Exception:
                    if attempt == RETRY_MAX - 1:
                        raise

                delay = RETRY_BASE_DELAY * (2**attempt)
                await asyncio.sleep(delay)

    @staticmethod
    def _resolve_account(ctx):
        from yuxi.channel.extensions.wecom.config import WeComConfig

        config = WeComConfig()
        return config.resolve_account()