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

from __future__ import annotations

import asyncio
import json
import logging

import httpx

from yuxi.channel.extensions.dingtalk.media import upload_media

logger = logging.getLogger(__name__)

SINGLE_API = "https://api.dingtalk.com/v1.0/robot/oToMessages/batchSend"
GROUP_API = "https://api.dingtalk.com/v1.0/robot/groupMessages/send"
RECALL_OTO_URL = "https://api.dingtalk.com/v1.0/robot/otoMessages/batchRecall"
RECALL_GROUP_URL = "https://api.dingtalk.com/v1.0/robot/groupMessages/recall"
QUERY_OTO_URL = "https://api.dingtalk.com/v1.0/robot/otoMessages/query"
QUERY_GROUP_URL = "https://api.dingtalk.com/v1.0/robot/groupMessages/query"
PLUGIN_SET_URL = "https://api.dingtalk.com/v1.0/robot/plugin/set"
CREATE_GROUP_URL = "https://api.dingtalk.com/v1.0/im/chat/scenegroup/create"
GROUP_MEMBERS_ADD_URL = "https://api.dingtalk.com/v1.0/im/sceneGroup/members/batchAdd"
GROUP_MEMBERS_REMOVE_URL = "https://api.dingtalk.com/v1.0/im/sceneGroup/members/batchRemove"
SEND_DING_URL = "https://api.dingtalk.com/v1.0/robot/ding/send"
BOT_LIST_IN_GROUP_URL = "https://api.dingtalk.com/v1.0/im/sceneGroup/robots/query"
BOT_GROUP_INFO_URL = "https://api.dingtalk.com/v1.0/robot/groupInfos/query"
UPDATE_ROBOT_URL = "https://api.dingtalk.com/v1.0/robot/info/update"
EXECUTE_AI_SKILL_URL = "https://api.dingtalk.com/v1.0/robot/aiSkills/execute"


class DingTalkOutbound:
    def __init__(self, gateway):
        self._gateway = gateway
        self._http: httpx.AsyncClient | None = None
        self._card_manager = None

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

    def set_card_manager(self, card_manager) -> None:
        self._card_manager = card_manager

    async def _ensure_http(self) -> httpx.AsyncClient:
        if self._http is None:
            self._http = httpx.AsyncClient(timeout=30.0)
        return self._http

    async def _get_token(self) -> str | None:
        return await self._gateway.token_manager.get_access_token()

    async def send_text(self, target_id: str, content: str, **kwargs) -> dict:
        reply_to_id = kwargs.pop("reply_to_id", None)
        msg_content = content
        if reply_to_id:
            msg_content = f"> 回复消息\n\n{content}"

        return await self._send_with_msg_key(
            target_id,
            msg_content,
            "sampleText",
            json.dumps({"content": msg_content}),
            **kwargs,
        )

    async def send_markdown(self, target_id: str, title: str, text: str, **kwargs) -> dict:
        reply_to_id = kwargs.pop("reply_to_id", None)
        msg_text = text
        if reply_to_id:
            msg_text = f"> 回复消息\n\n{text}"

        return await self._send_with_msg_key(
            target_id,
            msg_text,
            "sampleMarkdown",
            json.dumps({"title": title, "text": msg_text}),
            **kwargs,
        )

    async def send_link(
        self,
        target_id: str,
        title: str,
        text: str,
        message_url: str,
        pic_url: str = "",
        **kwargs,
    ) -> dict:
        msg_param = json.dumps(
            {
                "title": title,
                "text": text,
                "messageUrl": message_url,
                "picUrl": pic_url,
            }
        )
        return await self._send_with_msg_key(
            target_id,
            text,
            "sampleLink",
            msg_param,
            **kwargs,
        )

    async def send_image(self, target_id: str, image_path: str, **kwargs) -> dict:
        http = await self._ensure_http()
        token = await self._get_token()
        if not token:
            return {"success": False, "error": "no token"}

        media_id = await upload_media(http, token, image_path, "image")
        if not media_id:
            return {"success": False, "error": "upload failed"}

        text_content = kwargs.pop("text_content", None)
        if text_content:
            await self.send_text(target_id, text_content, **kwargs)
            await asyncio.sleep(0.3)

        return await self._send_with_msg_key(
            target_id,
            image_path,
            "sampleImageMsg",
            json.dumps({"photoURL": media_id}),
            **kwargs,
        )

    async def send_file(self, target_id: str, file_path: str, **kwargs) -> dict:
        http = await self._ensure_http()
        token = await self._get_token()
        if not token:
            return {"success": False, "error": "no token"}

        ext = file_path.rsplit(".", 1)[-1].lower() if "." in file_path else ""
        media_id = await upload_media(http, token, file_path, "file")
        if not media_id:
            return {"success": False, "error": "upload failed"}

        return await self._send_with_msg_key(
            target_id,
            file_path,
            "sampleFile",
            json.dumps(
                {
                    "mediaId": media_id,
                    "fileName": file_path.rsplit("/", 1)[-1].rsplit("\\", 1)[-1],
                    "fileType": ext,
                }
            ),
            **kwargs,
        )

    async def send_video(self, target_id: str, file_path: str, **kwargs) -> dict:
        http = await self._ensure_http()
        token = await self._get_token()
        if not token:
            return {"success": False, "error": "no token"}

        ext = file_path.rsplit(".", 1)[-1].lower() if "." in file_path else ""
        media_id = await upload_media(http, token, file_path, "video")
        if not media_id:
            return {"success": False, "error": "upload failed"}

        return await self._send_with_msg_key(
            target_id,
            file_path,
            "sampleVideo",
            json.dumps(
                {
                    "duration": "30",
                    "videoMediaId": media_id,
                    "videoType": ext,
                    "height": "400",
                    "width": "600",
                }
            ),
            **kwargs,
        )

    async def send_action_card(self, target_id: str, card: dict, **kwargs) -> dict:
        text = card.get("text", "")
        title = card.get("title", "通知")
        buttons = card.get("btns", [])
        single_title = card.get("singleTitle", "")
        single_url = card.get("singleURL", "")

        btn_count = len(buttons)

        if btn_count == 0 and single_title and single_url:
            msg_key = "sampleActionCard"
            msg_param = json.dumps(
                {
                    "title": title,
                    "text": text,
                    "singleTitle": single_title,
                    "singleURL": single_url,
                }
            )
        elif 1 <= btn_count <= 2:
            msg_key = "sampleActionCard2"
            msg_param = json.dumps(
                {
                    "title": title,
                    "text": text,
                    "btnOrientation": card.get("btnOrientation", "1"),
                    "actionButtons": buttons,
                }
            )
        elif btn_count == 3:
            msg_key = "sampleActionCard3"
            msg_param = json.dumps(
                {
                    "title": title,
                    "text": text,
                    "btnOrientation": card.get("btnOrientation", "1"),
                    "actionButtons": buttons,
                }
            )
        elif btn_count == 4:
            msg_key = "sampleActionCard4"
            msg_param = json.dumps(
                {
                    "title": title,
                    "text": text,
                    "btnOrientation": card.get("btnOrientation", "1"),
                    "actionButtons": buttons,
                }
            )
        else:
            return {"success": False, "error": f"unsupported button count: {btn_count}"}

        return await self._send_with_msg_key(
            target_id,
            text,
            msg_key,
            msg_param,
            **kwargs,
        )

    async def send_media(self, target_id: str, media_url: str, media_type: str, **kwargs) -> dict:
        if media_type.startswith("image"):
            return await self.send_image(target_id, media_url, **kwargs)
        elif media_type.startswith("video"):
            return await self.send_video(target_id, media_url, **kwargs)
        elif media_type.startswith("file"):
            return await self.send_file(target_id, media_url, **kwargs)
        return {"success": False, "error": f"unsupported media type: {media_type}"}

    async def recall_message(
        self,
        target_id: str,
        process_query_keys: list[str],
        *,
        is_group: bool = False,
    ) -> dict:
        http = await self._ensure_http()
        token = await self._get_token()
        if not token:
            return {"success": False, "error": "no token"}

        robot_code = self._gateway.robot_code

        if is_group:
            url = RECALL_GROUP_URL
            body = {
                "robotCode": robot_code,
                "openConversationId": target_id,
                "processQueryKeys": process_query_keys,
            }
        else:
            url = RECALL_OTO_URL
            body = {
                "robotCode": robot_code,
                "processQueryKeys": process_query_keys,
            }

        try:
            resp = await http.post(
                url,
                headers={
                    "x-acs-dingtalk-access-token": token,
                    "Content-Type": "application/json",
                },
                json=body,
                timeout=10.0,
            )
            data = resp.json()
            return {"success": True, "response": data}
        except Exception:
            logger.exception("DingTalk recall failed")
            return {"success": False, "error": "recall failed"}

    async def query_message_status(
        self,
        process_query_key: str,
        *,
        is_group: bool = False,
        open_conversation_id: str = "",
    ) -> dict:
        http = await self._ensure_http()
        token = await self._get_token()
        if not token:
            return {"success": False, "error": "no token"}

        robot_code = self._gateway.robot_code

        if is_group:
            url = QUERY_GROUP_URL
            body = {
                "robotCode": robot_code,
                "processQueryKey": process_query_key,
                "openConversationId": open_conversation_id,
            }
        else:
            url = QUERY_OTO_URL
            body = {
                "robotCode": robot_code,
                "processQueryKey": process_query_key,
            }

        try:
            resp = await http.post(
                url,
                headers={
                    "x-acs-dingtalk-access-token": token,
                    "Content-Type": "application/json",
                },
                json=body,
                timeout=10.0,
            )
            data = resp.json()
            return {"success": True, "response": data}
        except Exception:
            logger.exception("DingTalk query message status failed")
            return {"success": False, "error": "query failed"}

    async def set_robot_plugin(self, plugin_config: dict) -> dict:
        http = await self._ensure_http()
        token = await self._get_token()
        if not token:
            return {"success": False, "error": "no token"}

        robot_code = self._gateway.robot_code
        body = {"robotCode": robot_code, **plugin_config}

        try:
            resp = await http.post(
                PLUGIN_SET_URL,
                headers={
                    "x-acs-dingtalk-access-token": token,
                    "Content-Type": "application/json",
                },
                json=body,
                timeout=10.0,
            )
            data = resp.json()
            return {"success": True, "response": data}
        except Exception:
            logger.exception("DingTalk set robot plugin failed")
            return {"success": False, "error": "plugin set failed"}

    async def create_group(
        self,
        title: str,
        owner_user_id: str,
        *,
        user_ids: list[str] | None = None,
    ) -> dict:
        http = await self._ensure_http()
        token = await self._get_token()
        if not token:
            return {"success": False, "error": "no token"}

        body = {
            "title": title,
            "ownerUserId": owner_user_id,
        }
        if user_ids:
            body["userIds"] = user_ids

        try:
            resp = await http.post(
                CREATE_GROUP_URL,
                headers={
                    "x-acs-dingtalk-access-token": token,
                    "Content-Type": "application/json",
                },
                json=body,
                timeout=10.0,
            )
            data = resp.json()
            return {"success": True, "response": data}
        except Exception:
            logger.exception("DingTalk create group failed")
            return {"success": False, "error": "create failed"}

    async def add_group_members(self, open_conversation_id: str, user_ids: list[str]) -> dict:
        http = await self._ensure_http()
        token = await self._get_token()
        if not token:
            return {"success": False, "error": "no token"}

        try:
            resp = await http.post(
                GROUP_MEMBERS_ADD_URL,
                headers={
                    "x-acs-dingtalk-access-token": token,
                    "Content-Type": "application/json",
                },
                json={
                    "openConversationId": open_conversation_id,
                    "userIds": user_ids,
                },
                timeout=10.0,
            )
            data = resp.json()
            return {"success": True, "response": data}
        except Exception:
            logger.exception("DingTalk add group members failed")
            return {"success": False, "error": "add members failed"}

    async def remove_group_members(self, open_conversation_id: str, user_ids: list[str]) -> dict:
        http = await self._ensure_http()
        token = await self._get_token()
        if not token:
            return {"success": False, "error": "no token"}

        try:
            resp = await http.post(
                GROUP_MEMBERS_REMOVE_URL,
                headers={
                    "x-acs-dingtalk-access-token": token,
                    "Content-Type": "application/json",
                },
                json={
                    "openConversationId": open_conversation_id,
                    "userIds": user_ids,
                },
                timeout=10.0,
            )
            data = resp.json()
            return {"success": True, "response": data}
        except Exception:
            logger.exception("DingTalk remove group members failed")
            return {"success": False, "error": "remove members failed"}

    async def send_ding(
        self,
        user_ids: list[str],
        content: str,
        *,
        ding_type: str = "APP",
    ) -> dict:
        http = await self._ensure_http()
        token = await self._get_token()
        if not token:
            return {"success": False, "error": "no token"}

        robot_code = self._gateway.robot_code
        body = {
            "robotCode": robot_code,
            "userIds": user_ids,
            "content": content,
            "type": ding_type,
        }

        try:
            resp = await http.post(
                SEND_DING_URL,
                headers={
                    "x-acs-dingtalk-access-token": token,
                    "Content-Type": "application/json",
                },
                json=body,
                timeout=10.0,
            )
            data = resp.json()
            return {"success": True, "response": data}
        except Exception:
            logger.exception("DingTalk send DING failed")
            return {"success": False, "error": "send failed"}

    async def get_bot_list_in_group(self, open_conversation_id: str) -> dict:
        http = await self._ensure_http()
        token = await self._get_token()
        if not token:
            return {"success": False, "error": "no token"}

        try:
            resp = await http.post(
                BOT_LIST_IN_GROUP_URL,
                headers={
                    "x-acs-dingtalk-access-token": token,
                    "Content-Type": "application/json",
                },
                json={"openConversationId": open_conversation_id},
                timeout=10.0,
            )
            data = resp.json()
            return {"success": True, "response": data}
        except Exception:
            logger.exception("DingTalk get bot list failed")
            return {"success": False, "error": "query failed"}

    async def query_bot_group_info(self, open_conversation_id: str) -> dict:
        http = await self._ensure_http()
        token = await self._get_token()
        if not token:
            return {"success": False, "error": "no token"}

        robot_code = self._gateway.robot_code
        try:
            resp = await http.post(
                BOT_GROUP_INFO_URL,
                headers={
                    "x-acs-dingtalk-access-token": token,
                    "Content-Type": "application/json",
                },
                json={
                    "robotCode": robot_code,
                    "openConversationId": open_conversation_id,
                },
                timeout=10.0,
            )
            data = resp.json()
            return {"success": True, "response": data}
        except Exception:
            logger.exception("DingTalk query bot group info failed")
            return {"success": False, "error": "query failed"}

    async def update_robot_info(self, robot_info: dict) -> dict:
        http = await self._ensure_http()
        token = await self._get_token()
        if not token:
            return {"success": False, "error": "no token"}

        robot_code = self._gateway.robot_code
        body = {"robotCode": robot_code, **robot_info}

        try:
            resp = await http.post(
                UPDATE_ROBOT_URL,
                headers={
                    "x-acs-dingtalk-access-token": token,
                    "Content-Type": "application/json",
                },
                json=body,
                timeout=10.0,
            )
            data = resp.json()
            return {"success": True, "response": data}
        except Exception:
            logger.exception("DingTalk update robot info failed")
            return {"success": False, "error": "update failed"}

    async def execute_ai_skill(
        self,
        skill_id: str,
        prompt: str,
        *,
        user_id: str = "",
        conversation_id: str = "",
    ) -> dict:
        http = await self._ensure_http()
        token = await self._get_token()
        if not token:
            return {"success": False, "error": "no token"}

        robot_code = self._gateway.robot_code
        body = {
            "robotCode": robot_code,
            "skillId": skill_id,
            "prompt": prompt,
        }
        if user_id:
            body["userId"] = user_id
        if conversation_id:
            body["openConversationId"] = conversation_id

        try:
            resp = await http.post(
                EXECUTE_AI_SKILL_URL,
                headers={
                    "x-acs-dingtalk-access-token": token,
                    "Content-Type": "application/json",
                },
                json=body,
                timeout=30.0,
            )
            data = resp.json()
            return {"success": True, "response": data}
        except Exception:
            logger.exception("DingTalk execute AI skill failed")
            return {"success": False, "error": "execute failed"}

    async def streaming_send(
        self,
        target_id: str,
        content: str,
        *,
        is_group: bool = False,
        conversation_id: str = "",
        status: str = "PROCESSING",
    ) -> dict:
        if not self._card_manager:
            return {"success": False, "error": "card manager not initialized"}

        success = await self._card_manager.streaming_update(content, status)
        return {"success": success}

    async def _send_with_msg_key(
        self,
        target_id: str,
        content: str,
        msg_key: str,
        msg_param: str,
        **kwargs,
    ) -> dict:
        http = await self._ensure_http()
        token = await self._get_token()
        if not token:
            return {"success": False, "error": "no token"}

        is_group = kwargs.get("is_group", False)
        conversation_id = kwargs.get("conversation_id", "")
        sender_staff_id = kwargs.get("sender_staff_id", "")

        robot_code = self._gateway.robot_code

        if is_group:
            url = GROUP_API
            body = {
                "robotCode": robot_code,
                "openConversationId": conversation_id or target_id,
                "msgKey": msg_key,
                "msgParam": msg_param,
            }
        else:
            url = SINGLE_API
            user_ids = kwargs.get("user_ids")
            if not user_ids and sender_staff_id:
                user_ids = [sender_staff_id]
            if not user_ids:
                user_ids = [target_id]

            body = {
                "robotCode": robot_code,
                "userIds": user_ids,
                "msgKey": msg_key,
                "msgParam": msg_param,
            }

        try:
            resp = await http.post(
                url,
                headers={
                    "x-acs-dingtalk-access-token": token,
                    "Content-Type": "application/json",
                },
                json=body,
                timeout=30.0,
            )
            data = resp.json()
            logger.debug("DingTalk send result: %s", data)
            return {"success": True, "response": data}
        except Exception:
            logger.exception("DingTalk send failed")
            return {"success": False, "error": "send failed"}