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

from __future__ import annotations

import asyncio
import base64
import logging
import random
import struct
import time
from typing import Any

import httpx

from yuxi.channel.extensions.wechat_ilink.errors import ILinkApiError, ILinkErrorCode, map_api_error
from yuxi.channel.extensions.wechat_ilink.types import ILinkAccount, ILinkMessage

logger = logging.getLogger(__name__)

BASE_URL = "https://ilinkai.weixin.qq.com"
BASE_INFO = {"channel_version": "2.0.0"}

BACKOFF_SEQUENCE = [1.0, 2.0, 5.0, 10.0, 30.0, 60.0]
MAX_RETRIES = 100
FAST_DISCONNECT_WINDOW_S = 5.0
FAST_DISCONNECT_THRESHOLD = 3
FAST_DISCONNECT_PENALTY_S = 60.0


class RetryController:
    def __init__(self, max_retries: int = MAX_RETRIES):
        self._max_retries = max_retries
        self._attempt = 0
        self._disconnect_times: list[float] = []

    def reset(self) -> None:
        self._attempt = 0
        self._disconnect_times.clear()

    def record_disconnect(self) -> None:
        now = time.time()
        self._disconnect_times.append(now)
        cutoff = now - FAST_DISCONNECT_WINDOW_S
        self._disconnect_times = [t for t in self._disconnect_times if t > cutoff]

    @property
    def fast_disconnect(self) -> bool:
        return len(self._disconnect_times) >= FAST_DISCONNECT_THRESHOLD

    @property
    def exhausted(self) -> bool:
        return self._attempt >= self._max_retries

    def next_delay(self) -> float:
        if self.fast_disconnect:
            return FAST_DISCONNECT_PENALTY_S

        idx = min(self._attempt, len(BACKOFF_SEQUENCE) - 1)
        delay = BACKOFF_SEQUENCE[idx]
        self._attempt += 1
        return delay

    async def wait(self) -> None:
        delay = self.next_delay()
        logger.info("iLink poll 重连等待 %.1fs (attempt #%d)", delay, self._attempt)
        await asyncio.sleep(delay)


class ILinkApiClient:
    def __init__(self, account: ILinkAccount):
        self._account = account
        self._client: httpx.AsyncClient | None = None

    async def __aenter__(self):
        self._client = httpx.AsyncClient(
            base_url=self._account.base_url,
            timeout=httpx.Timeout(35.0),
            proxy=self._account.proxy,
        )
        return self

    async def __aexit__(self, *args):
        if self._client:
            await self._client.aclose()

    def _auth_headers(self) -> dict[str, str]:
        uin_bytes = struct.pack(">I", random.randint(0, 0xFFFFFFFF))
        x_wechat_uin = base64.b64encode(uin_bytes).decode("ascii")
        return {
            "Authorization": f"Bearer {self._account.bot_token}",
            "AuthorizationType": "ilink_bot_token",
            "X-WECHAT-UIN": x_wechat_uin,
        }

    async def get_qrcode(self) -> dict[str, Any]:
        resp = await self._client.get("/ilink/bot/get_bot_qrcode", params={"bot_type": 3})
        return self._check_response(resp)

    async def check_login_status(self, qrcode_id: str) -> dict[str, Any]:
        resp = await self._client.get("/ilink/bot/get_qrcode_status", params={"qrcode": qrcode_id})
        data = self._check_response(resp)

        status = data.get("status", "")
        if status == "confirmed":
            bot_token = data.get("bot_token", "")
            if bot_token:
                self._account.bot_token = bot_token
            ilink_bot_id = data.get("ilink_bot_id", "")
            if ilink_bot_id:
                self._account.ilink_bot_id = ilink_bot_id
            ilink_user_id = data.get("ilink_user_id", "")
            if ilink_user_id:
                self._account.ilink_user_id = ilink_user_id
            baseurl = data.get("baseurl", "")
            if baseurl:
                self._account.base_url = baseurl
                if self._client:
                    await self._client.aclose()
                self._client = httpx.AsyncClient(
                    base_url=self._account.base_url,
                    timeout=httpx.Timeout(35.0),
                    proxy=self._account.proxy,
                )

        return data

    async def poll_messages(
        self,
        timeout: int = 25,
        updates_buf: str | None = None,
    ) -> tuple[list[dict[str, Any]], str | None]:
        body = {"timeout": timeout, "base_info": BASE_INFO}
        if updates_buf:
            body["get_updates_buf"] = updates_buf
        resp = await self._client.post(
            "/ilink/bot/getupdates",
            headers=self._auth_headers(),
            json=body,
        )
        data = self._check_response(resp)
        return data.get("updates", []), data.get("get_updates_buf")

    async def send_message(self, payload: dict[str, Any]) -> dict[str, Any]:
        if "base_info" not in payload:
            payload["base_info"] = BASE_INFO
        resp = await self._client.post(
            "/ilink/bot/sendmessage",
            headers=self._auth_headers(),
            json=payload,
        )
        return self._check_response(resp)

    async def get_upload_url(self, media_type: str, filesize: int) -> dict[str, Any]:
        body = {
            "media_type": media_type,
            "filesize": filesize,
            "base_info": BASE_INFO,
        }
        resp = await self._client.post(
            "/ilink/bot/getuploadurl",
            headers=self._auth_headers(),
            json=body,
        )
        return self._check_response(resp)

    async def get_config(self, ilink_user_id: str, context_token: str) -> dict[str, Any]:
        body = {
            "ilink_user_id": ilink_user_id,
            "context_token": context_token,
            "base_info": BASE_INFO,
        }
        resp = await self._client.post(
            "/ilink/bot/getconfig",
            headers=self._auth_headers(),
            json=body,
        )
        return self._check_response(resp)

    async def send_typing(self, user_id: str, typing_ticket: str, status: int) -> dict[str, Any]:
        body = {
            "msg": {
                "to_user_id": user_id,
                "from_user_id": "",
                "message_type": 2,
                "message_state": 2,
                "typing_ticket": typing_ticket,
                "status": status,
            },
            "base_info": BASE_INFO,
        }
        resp = await self._client.post(
            "/ilink/bot/sendtyping",
            headers=self._auth_headers(),
            json=body,
        )
        return self._check_response(resp)

    @staticmethod
    def _check_response(resp: httpx.Response) -> dict[str, Any]:
        if resp.is_success:
            data = resp.json()
            errcode = data.get("errcode", 0)
            if errcode != 0:
                error = map_api_error(errcode, data.get("errmsg", ""))
                if error:
                    raise error
            return data
        resp.raise_for_status()
        return {}


class WeChatILinkGateway:
    def __init__(
        self,
        account: ILinkAccount,
        message_handler,
        status_adapter=None,
        context_store=None,
    ):
        self._account = account
        self._message_handler = message_handler
        self._status = status_adapter
        self._context_store = context_store

        self._poll_task: asyncio.Task | None = None
        self._running = False
        self._retry = RetryController(max_retries=account.max_retries)
        self._api: ILinkApiClient | None = None
        self._updates_buf: str | None = None

    async def start(self, ctx) -> None:
        if not self._account.token_valid:
            raise ILinkApiError(
                ILinkErrorCode.AUTH_FAILED,
                "bot_token 无效或已过期，请重新扫码登录",
            )

        self._running = True
        self._api = ILinkApiClient(self._account)
        await self._api.__aenter__()

        self._poll_task = asyncio.create_task(self._poll_loop(ctx), name="ilink-poll")

        if self._status:
            self._status.connected = True

        logger.info("iLink Gateway 已启动 (account=%s)", self._account.account_id)

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

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

        if self._api:
            await self._api.__aexit__(None, None, None)
            self._api = None

        if self._status:
            self._status.connected = False

        logger.info("iLink Gateway 已停止 (account=%s)", self._account.account_id)

    async def _poll_loop(self, ctx) -> None:
        while self._running and not ctx.cancel_event.is_set():
            try:
                raw_messages, self._updates_buf = await self._api.poll_messages(
                    timeout=self._account.poll_timeout,
                    updates_buf=self._updates_buf,
                )
                self._retry.reset()

                if self._status:
                    self._status.record_poll_success()

                for raw in raw_messages:
                    try:
                        msg = ILinkMessage.from_api_response(raw)
                        if self._message_handler:
                            await self._message_handler(msg)
                    except Exception:
                        logger.exception("iLink 消息处理异常: %s", raw.get("msg_id"))

            except ILinkApiError as e:
                if e.errcode == -14:
                    logger.error("iLink 会话过期 (errcode=-14)，触发恢复流程")
                    if self._context_store:
                        self._context_store.clear_all()
                    self._updates_buf = None
                    self._account.bot_token = ""
                    break

                if e.code in (ILinkErrorCode.INVALID_TOKEN, ILinkErrorCode.TOKEN_EXPIRED):
                    logger.error("iLink token 失效，等待刷新后重试")
                    if self._status:
                        self._status.record_token_error()
                    await asyncio.sleep(5.0)
                    continue

                if self._status:
                    self._status.record_poll_error()
                self._retry.record_disconnect()
                await self._retry.wait()

            except (httpx.TimeoutException, httpx.NetworkError):
                if self._status:
                    self._status.record_poll_error()
                self._retry.record_disconnect()
                await self._retry.wait()

            except httpx.HTTPStatusError as e:
                logger.error("iLink HTTP 错误: %s", e)
                if self._status:
                    self._status.record_poll_error()
                self._retry.record_disconnect()
                await self._retry.wait()

            except asyncio.CancelledError:
                break

            except Exception:
                logger.exception("iLink poll 循环异常")
                if self._status:
                    self._status.record_poll_error()
                self._retry.record_disconnect()
                await self._retry.wait()

            finally:
                await asyncio.sleep(self._account.poll_interval)

        logger.info("iLink poll loop 退出 (account=%s)", self._account.account_id)