ForcePilot/backend/package/yuxi/channel/extensions/wechat-mp/gateway.py

import asyncio
import logging
import time

import httpx

from yuxi.channel.extensions.wechat_mp.config import WeChatMPConfig

logger = logging.getLogger(__name__)

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

_queue: asyncio.Queue | None = None
_deduplicator = None
_gateway: "WeChatMPGateway | None" = None


def _get_webhook_queue() -> asyncio.Queue | None:
    return _queue


def _get_deduplicator():
    global _deduplicator
    if _deduplicator is None:
        from yuxi.channel.extensions.wechat_mp.dedupe import MessageDeduplicator

        _deduplicator = MessageDeduplicator()
    return _deduplicator


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


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


class WeChatMPGateway:

    def __init__(self):
        self._account = None
        self._queue: asyncio.Queue | None = 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:
        global _queue
        account = self._resolve_account(ctx)
        if not account.is_configured():
            logger.warning("WeChat MP account not configured, skipping start")
            return {"running": False, "reason": "not-configured"}

        self._account = account
        _queue = asyncio.Queue(maxsize=1000)
        self._queue = _queue
        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("WeChat MP gateway started for account %s", account.account_id)
            return {"running": True, "account_id": account.account_id}
        except Exception:
            logger.exception("WeChat MP gateway failed to start")
            await self.stop(ctx)
            raise

    async def stop(self, ctx) -> None:
        global _queue
        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

        _queue = None
        self._queue = None
        logger.info("WeChat MP 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

    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

            params = {
                "grant_type": "client_credential",
                "appid": self._account.app_id,
                "secret": self._account.app_secret,
            }

            for attempt in range(RETRY_MAX):
                try:
                    resp = await self._http.get(TOKEN_URL, params=params)
                    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("WeChat MP 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):
        config = WeChatMPConfig()
        return config.resolve_account()