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

from __future__ import annotations

import asyncio
import json
import logging
import time
from typing import TYPE_CHECKING

import httpx

from yuxi.channel.extensions.dingtalk.types import DingTalkAccount

if TYPE_CHECKING:
    from yuxi.channel.extensions.dingtalk.monitor import DingTalkMonitor

logger = logging.getLogger(__name__)

TOKEN_URL = "https://api.dingtalk.com/v1.0/oauth2/accessToken"
TOKEN_REFRESH_MARGIN = 300
RECONNECT_WAIT_MS = 100
RECONNECT_MAX_WAITS = 100
OPEN_CONNECTION_URL = "https://api.dingtalk.com/v1.0/gateway/connections/open"


class DingTalkTokenManager:
    def __init__(self, app_key: str, app_secret: str):
        self._app_key = app_key
        self._app_secret = app_secret
        self._token: str | None = None
        self._expires_at: float = 0.0
        self._lock = asyncio.Lock()
        self._http: httpx.AsyncClient | None = None

    def set_http(self, http: httpx.AsyncClient) -> None:
        self._http = http

    @property
    def token(self) -> str | None:
        if self._token and time.time() < self._expires_at:
            return self._token
        return None

    async def get_access_token(self) -> str | None:
        if self._token and time.time() < self._expires_at:
            return self._token

        async with self._lock:
            if self._token and time.time() < self._expires_at:
                return self._token

            try:
                resp = await self._http.post(
                    TOKEN_URL,
                    json={
                        "appKey": self._app_key,
                        "appSecret": self._app_secret,
                    },
                    timeout=10.0,
                )
                data = resp.json()
                self._token = data["accessToken"]
                expire_in = data.get("expireIn", 7200)
                self._expires_at = time.time() + expire_in - TOKEN_REFRESH_MARGIN
                logger.info("DingTalk access_token refreshed, expires_in=%ds", expire_in)
                return self._token
            except Exception:
                logger.exception("Failed to get DingTalk access_token")
                return None


class DingTalkGateway:
    def __init__(self, account: DingTalkAccount, token_manager: DingTalkTokenManager):
        self._account = account
        self._token_manager = token_manager
        self._running = False
        self._event_loop: asyncio.AbstractEventLoop | None = None
        self._robot_code_cache: str | None = None
        self._handler: DingTalkMonitor | None = None
        self._stream_client = None
        self._credential = None

    def set_handler(self, handler: DingTalkMonitor) -> None:
        self._handler = handler

    @property
    def robot_code(self) -> str | None:
        return self._robot_code_cache or self._account.robot_code

    @property
    def app_key(self) -> str:
        return self._account.app_key

    @property
    def app_secret(self) -> str:
        return self._account.app_secret

    @property
    def token_manager(self) -> DingTalkTokenManager:
        return self._token_manager

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

        self._running = True
        self._event_loop = asyncio.get_running_loop()

        try:
            from dingtalk_stream import ChatbotMessage, Credential, DingTalkStreamClient

            self._credential = Credential(self._account.app_key, self._account.app_secret)

            self._stream_client = DingTalkStreamClient(
                credential=self._credential,
            )
            self._stream_client.register_callback_handler(ChatbotMessage.TOPIC, self._handler)

            extra_topics = getattr(self._account, "subscribe_topics", [])
            for topic in extra_topics:
                try:
                    self._stream_client.register_callback_handler(topic, self._handler)
                    logger.info("DingTalk extra topic registered: %s", topic)
                except Exception:
                    logger.warning("Failed to register topic: %s", topic, exc_info=True)

            try:
                self._stream_client.register_callback_handler("/v1.0/card/instances/callback", self._handler)
                logger.info("DingTalk card callback registered")
            except Exception:
                logger.debug("Failed to register card callback", exc_info=True)

            self._stream_client.pre_start()
        except ImportError:
            logger.error("dingtalk-stream-sdk-python not installed")
            return {"running": False, "reason": "sdk-missing"}
        except Exception:
            logger.exception("Failed to initialize DingTalk stream client")
            return {"running": False, "reason": "init-failed"}

        asyncio.create_task(self._connection_loop(), name="dingtalk-stream-loop")
        logger.info("DingTalk gateway started for account %s", self._account.account_id)
        return {"running": True, "account_id": self._account.account_id}

    async def stop(self, ctx) -> None:
        self._running = False
        self._stream_client = None
        logger.info("DingTalk gateway stopped")

    async def _open_connection(self) -> tuple[str, str]:
        async with httpx.AsyncClient(timeout=10.0) as http:
            resp = await http.post(
                OPEN_CONNECTION_URL,
                headers={
                    "Content-Type": "application/json",
                    "clientId": self._account.app_key,
                    "dingtalk-accept-encoding": "base64",
                },
            )
            resp.raise_for_status()
            data = resp.json()
            return data["endpoint"], data["ticket"]

    async def _connection_loop(self):
        first_attempt = True

        while self._running:
            try:
                endpoint, ticket = await self._open_connection()
            except Exception as e:
                if first_attempt:
                    logger.error("DingTalk first connection failed: %s", e)
                first_attempt = False
                await self._wait_with_check(RECONNECT_MAX_WAITS)
                continue

            try:
                import websockets

                async with websockets.connect(f"{endpoint}?ticket={ticket}") as ws:
                    if first_attempt:
                        logger.info("DingTalk stream connected")
                        first_attempt = False

                    async for raw_message in ws:
                        if not self._running:
                            break
                        try:
                            json_message = json.loads(raw_message)
                            result = self._stream_client.route_message(json_message)
                            if result == "TAG_DISCONNECT":
                                logger.warning("DingTalk stream: TAG_DISCONNECT received")
                                break
                        except Exception:
                            logger.exception("DingTalk stream message routing error")
            except Exception:
                logger.exception("DingTalk stream connection error")
                await self._wait_with_check(RECONNECT_MAX_WAITS)

    async def _wait_with_check(self, max_waits: int):
        for _ in range(max_waits):
            if not self._running:
                break
            await asyncio.sleep(RECONNECT_WAIT_MS / 1000)

    async def probe(self, account: dict) -> bool:
        return account.get("app_key") and account.get("app_secret")