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

from __future__ import annotations

import asyncio
import json
import logging
import time
from datetime import UTC, datetime

try:
    from dingtalk_stream import AckMessage, ChatbotMessage
except ImportError:
    AckMessage = None
    ChatbotMessage = None

from yuxi.channel.extensions.dingtalk.dedupe import DingTalkDeduplicator
from yuxi.channel.extensions.dingtalk.file_cache import FileCache

logger = logging.getLogger(__name__)

STALE_MSG_AGE_S = 60

BUILTIN_COMMANDS = {
    "/help": "查看帮助信息",
    "/status": "查看 Bot 运行状态",
    "/ping": "测试连通性",
}


class CommandResult:
    __slots__ = ("handled", "response")

    def __init__(self, handled: bool = False, response: str = ""):
        self.handled = handled
        self.response = response


class DingTalkMonitor:
    def __init__(self):
        self._robot_code_cache: str | None = None
        self._deduplicator: DingTalkDeduplicator | None = None
        self._file_cache: FileCache | None = None
        self._security = None
        self._pairing = None
        self._outbound = None
        self._gateway = None
        self._dispatch_callback = None

    def inject_dependencies(
        self,
        *,
        gateway=None,
        outbound=None,
        security=None,
        pairing=None,
        deduplicator: DingTalkDeduplicator | None = None,
        file_cache: FileCache | None = None,
    ) -> None:
        self._gateway = gateway
        self._outbound = outbound
        self._security = security
        self._pairing = pairing
        self._deduplicator = deduplicator
        self._file_cache = file_cache

    def set_dispatch_callback(self, callback) -> None:
        self._dispatch_callback = callback

    async def process(self, callback):
        topic = getattr(callback, "topic", "")
        if topic and "/card/instances/callback" in topic:
            return await self._handle_card_callback(callback)

        try:
            incoming = ChatbotMessage.from_dict(callback.data)
        except Exception:
            logger.exception("Failed to parse DingTalk ChatbotMessage")
            return AckMessage.STATUS_SYSTEM_EXCEPTION, "ERROR"

        self._robot_code_cache = incoming.robot_code

        msg_age_s = (time.time() * 1000 - incoming.create_at) / 1000
        if msg_age_s > STALE_MSG_AGE_S:
            logger.debug("DingTalk stale message ignored, age=%.1fs", msg_age_s)
            return AckMessage.STATUS_OK, "OK"

        is_group = incoming.conversation_type != "1"

        if self._deduplicator:
            create_time_s = incoming.create_at / 1000 if incoming.create_at else None
            if self._deduplicator.is_duplicate(incoming.message_id, create_time_s):
                logger.debug("Duplicate DingTalk message: %s", incoming.message_id)
                return AckMessage.STATUS_OK, "OK"

        if not is_group and self._deduplicator and self._deduplicator.is_my_msg(incoming):
            logger.debug("Filtered own message in direct chat: %s", incoming.message_id)
            return AckMessage.STATUS_OK, "OK"

        if incoming.message_type == "picture":
            image_path = await self._handle_picture(incoming, is_group)
            if image_path:
                content = f"[图片: {image_path}]"
                unified = self._build_unified_message(incoming, content, "IMAGE", is_group)
                if self._dispatch_callback:
                    asyncio.create_task(
                        self._dispatch_callback(unified),
                        name=f"dingtalk-dispatch-{incoming.message_id[:8]}",
                    )
            return AckMessage.STATUS_OK, "OK"

        content, msg_type = await self._parse_message_content(incoming, is_group)
        if content is None:
            return AckMessage.STATUS_OK, "OK"

        if msg_type == "TEXT" and content:
            cmd_result = await self._handle_command(content, incoming, is_group)
            if cmd_result.handled:
                return AckMessage.STATUS_OK, "OK"

        if is_group:
            from_user_id = incoming.conversation_id
            other_user_id = incoming.conversation_id
            actual_user_id = incoming.sender_id
        else:
            from_user_id = incoming.sender_id
            other_user_id = incoming.sender_id
            actual_user_id = incoming.sender_id

        unified = {
            "msg_id": incoming.message_id,
            "channel_type": "dingtalk",
            "account_id": "default",
            "content": content,
            "message_type": msg_type,
            "sender": {
                "id": from_user_id,
                "display_name": getattr(incoming, "sender_nick", None) or from_user_id,
                "kind": "GROUP" if is_group else "DIRECT",
            },
            "group": {
                "id": incoming.conversation_id,
                "type": "group" if is_group else "direct",
            }
            if is_group
            else None,
            "timestamp": datetime.fromtimestamp(incoming.create_at / 1000, tz=UTC),
            "mentioned_ids": self._extract_mentioned_ids(incoming),
            "reply_to_id": None,
            "thread_id": None,
            "session_key": incoming.conversation_id if is_group else incoming.sender_id,
            "raw_payload": incoming,
            "metadata": {
                "conversation_id": incoming.conversation_id,
                "conversation_type": "2" if is_group else "1",
                "sender_staff_id": getattr(incoming, "sender_staff_id", ""),
                "robot_code": self._robot_code_cache,
                "is_group": is_group,
                "actual_user_id": actual_user_id,
                "other_user_id": other_user_id,
            },
        }

        if self._dispatch_callback:
            asyncio.create_task(
                self._dispatch_callback(unified),
                name=f"dingtalk-dispatch-{incoming.message_id[:8]}",
            )
        else:
            logger.warning("DingTalk: no dispatch callback set, message dropped: %s", incoming.message_id)

        return AckMessage.STATUS_OK, "OK"

    async def _handle_picture(self, incoming, is_group: bool) -> str | None:
        image_list = incoming.get_image_list() if hasattr(incoming, "get_image_list") else []
        if not image_list:
            return None
        download_code = image_list[0]
        image_path = await self._download_image(incoming, download_code)
        if image_path and self._file_cache:
            cache_key = incoming.conversation_id if is_group else incoming.sender_id
            self._file_cache.add(cache_key, image_path)
            logger.debug("DingTalk image cached: %s -> %s", download_code, image_path)
        return image_path

    async def _parse_message_content(self, incoming, is_group: bool) -> tuple[str | None, str | None]:
        msg_type = incoming.message_type

        if msg_type == "text":
            text = getattr(incoming, "text", None)
            content = text.content.strip() if text else ""

            if self._file_cache:
                cache_key = incoming.conversation_id if is_group else incoming.sender_id
                cached = self._file_cache.get(cache_key)
                if cached:
                    content = f"{content}\n[图片: {cached.file_path}]" if content else f"[图片: {cached.file_path}]"
                    self._file_cache.clear(cache_key)

            return content, "TEXT"

        elif msg_type == "audio":
            extensions = getattr(incoming, "extensions", {})
            recognition = extensions.get("content", {}).get("recognition", "")
            return recognition, "TEXT"

        elif msg_type == "richText":
            text_list = incoming.get_text_list() if hasattr(incoming, "get_text_list") else []
            image_list = incoming.get_image_list() if hasattr(incoming, "get_image_list") else []
            text_content = "".join(text_list).strip()
            image_paths = []
            for download_code in image_list:
                path = await self._download_image(incoming, download_code)
                if path:
                    image_paths.append(path)
            if image_paths:
                image_lines = "\n".join(f"[图片: {p}]" for p in image_paths)
                text_content = f"{text_content}\n{image_lines}" if text_content else image_lines
            return text_content or None, "TEXT"

        elif msg_type == "video":
            download_code = getattr(incoming, "download_code", None) or getattr(incoming, "video_download_code", None)
            if download_code:
                video_path = await self._download_image(incoming, download_code)
                if video_path:
                    return f"[视频: {video_path}]", "VIDEO"
            return "[视频消息]", "VIDEO"

        elif msg_type == "file":
            file_name = getattr(incoming, "file_name", "") or "未知文件"
            download_code = getattr(incoming, "download_code", None) or getattr(incoming, "file_download_code", None)
            if download_code:
                file_path = await self._download_image(incoming, download_code)
                if file_path:
                    return f"[文件: {file_name} -> {file_path}]", "FILE"
            return f"[文件: {file_name}]", "FILE"

        return None, None

    async def _download_image(self, incoming, download_code: str) -> str | None:
        if not self._gateway or not self._gateway.token_manager:
            return None

        http = self._gateway.token_manager._http
        if http is None:
            async with __import__("httpx").AsyncClient(timeout=30.0) as http:
                return await self._do_download(http, download_code)
        return await self._do_download(http, download_code)

    def _extract_mentioned_ids(self, incoming) -> list[str]:
        at_users = getattr(incoming, "at_users", None)
        if at_users:
            return [u.get("dingtalkId", "") for u in at_users if u.get("dingtalkId")]
        raw = getattr(incoming, "raw_payload", None)
        if raw and isinstance(raw, dict):
            at_users_raw = raw.get("atUsers", [])
            return [u.get("dingtalkId", "") for u in at_users_raw if u.get("dingtalkId")]
        return []

    def _build_unified_message(self, incoming, content: str, msg_type: str, is_group: bool) -> dict:
        if is_group:
            from_user_id = incoming.conversation_id
            other_user_id = incoming.conversation_id
            actual_user_id = incoming.sender_id
        else:
            from_user_id = incoming.sender_id
            other_user_id = incoming.sender_id
            actual_user_id = incoming.sender_id

        return {
            "msg_id": incoming.message_id,
            "channel_type": "dingtalk",
            "account_id": "default",
            "content": content,
            "message_type": msg_type,
            "sender": {
                "id": from_user_id,
                "display_name": getattr(incoming, "sender_nick", None) or from_user_id,
                "kind": "GROUP" if is_group else "DIRECT",
            },
            "group": {
                "id": incoming.conversation_id,
                "type": "group" if is_group else "direct",
            }
            if is_group
            else None,
            "timestamp": datetime.fromtimestamp(incoming.create_at / 1000, tz=UTC),
            "mentioned_ids": self._extract_mentioned_ids(incoming),
            "reply_to_id": None,
            "thread_id": None,
            "session_key": incoming.conversation_id if is_group else incoming.sender_id,
            "raw_payload": incoming,
            "metadata": {
                "conversation_id": incoming.conversation_id,
                "conversation_type": "2" if is_group else "1",
                "sender_staff_id": getattr(incoming, "sender_staff_id", ""),
                "robot_code": self._robot_code_cache,
                "is_group": is_group,
                "actual_user_id": actual_user_id,
                "other_user_id": other_user_id,
            },
        }

    async def _handle_card_callback(self, callback) -> tuple[str, str]:
        data = callback.data if hasattr(callback, "data") else {}
        if isinstance(data, str):
            try:
                data = json.loads(data)
            except Exception:
                data = {}

        card_instance_id = data.get("cardInstanceId", "")
        out_track_id = data.get("outTrackId", "")
        button_key = data.get("buttonKey", "")
        button_value = data.get("buttonValue", "")
        user_id = data.get("userId", "")
        conversation_id = data.get("conversationId", "")

        logger.info(
            "DingTalk card callback: card=%s button=%s value=%s user=%s",
            card_instance_id,
            button_key,
            button_value,
            user_id,
        )

        unified = {
            "msg_id": out_track_id or card_instance_id,
            "channel_type": "dingtalk",
            "account_id": "default",
            "content": f"[卡片回调] buttonKey={button_key} buttonValue={button_value}",
            "message_type": "CARD_CALLBACK",
            "sender": {"id": user_id, "display_name": user_id, "kind": "DIRECT"},
            "group": {"id": conversation_id, "type": "group"} if conversation_id else None,
            "timestamp": datetime.now(UTC),
            "mentioned_ids": [],
            "reply_to_id": None,
            "thread_id": None,
            "session_key": conversation_id or user_id,
            "raw_payload": data,
            "metadata": {
                "card_instance_id": card_instance_id,
                "out_track_id": out_track_id,
                "button_key": button_key,
                "button_value": button_value,
                "conversation_id": conversation_id,
                "user_id": user_id,
                "is_card_callback": True,
            },
        }

        if self._dispatch_callback:
            asyncio.create_task(
                self._dispatch_callback(unified),
                name=f"dingtalk-card-{out_track_id[:8] if out_track_id else 'unknown'}",
            )

        return AckMessage.STATUS_OK, "OK"

    async def _handle_command(self, text: str, incoming, is_group: bool) -> CommandResult:
        if not text.startswith("/"):
            return CommandResult()

        parts = text.strip().split(maxsplit=1)
        cmd = parts[0].lower()

        if cmd == "/help":
            help_lines = ["**可用命令**:"] + [f"- `{k}`: {v}" for k, v in BUILTIN_COMMANDS.items()]
            response = "\n".join(help_lines)
        elif cmd == "/ping":
            response = "pong 🏓"
        elif cmd == "/status":
            response = "运行正常 ✅"
        else:
            response = f"未知命令: `{cmd}`，输入 `/help` 查看帮助"

        if self._outbound:
            try:
                await self._outbound.send_text(
                    incoming.sender_id,
                    response,
                    is_group=is_group,
                    conversation_id=incoming.conversation_id if is_group else "",
                    sender_staff_id=getattr(incoming, "sender_staff_id", ""),
                )
            except Exception:
                logger.exception("DingTalk command reply failed")

        return CommandResult(handled=True, response=response)

    async def _do_download(self, http, download_code: str) -> str | None:
        token = await self._gateway.token_manager.get_access_token()
        if not token:
            return None

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

        robot_code = self._robot_code_cache or self._gateway.robot_code or ""
        return await download_image(http, token, download_code, robot_code)