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

from __future__ import annotations

import asyncio
import logging
from typing import Any

from yuxi.channel.extensions.qqbot.api_client import QQBotApiClient
from yuxi.channel.extensions.qqbot.format import chunk_markdown_text, md_to_qq_markdown
from yuxi.channel.extensions.qqbot.types import QQBotAccountConfig, QQBotChatType

logger = logging.getLogger(__name__)


class QQBotOutbound:
    def __init__(
        self,
        api_client: QQBotApiClient,
        account: QQBotAccountConfig,
        streaming_controller: Any | None = None,
    ):
        self._api_client = api_client
        self._account = account
        self._streaming_controller = streaming_controller

    async def send_text(
        self,
        target_id: str,
        content: str,
        reply_to_id: str | None = None,
        account_id: str | None = None,
    ) -> None:
        chat_type, target = self._parse_target(target_id)

        if not content:
            return

        message_reference = None
        if reply_to_id:
            message_reference = {"message_id": reply_to_id, "ignore_get_message_error": True}

        if self._account.markdown_support:
            formatted = md_to_qq_markdown(content)
            chunks = chunk_markdown_text(formatted)
            total = len(chunks)
            for i, chunk in enumerate(chunks):
                if total > 1:
                    chunk = f"{chunk}\n\n({i + 1}/{total})"
                if chat_type == QQBotChatType.C2C:
                    await self._api_client.send_c2c_markdown(
                        target, chunk,
                        msg_id=reply_to_id if i == 0 else None,
                        message_reference=message_reference if i == total - 1 else None,
                    )
                elif chat_type == QQBotChatType.GUILD:
                    await self._api_client.send_channel_message(
                        target, chunk, msg_type=0,
                        markdown={"content": chunk},
                    )
                else:
                    await self._api_client.send_group_markdown(
                        target, chunk,
                        msg_id=reply_to_id if i == 0 else None,
                        message_reference=message_reference if i == total - 1 else None,
                    )
                if i < total - 1:
                    await asyncio.sleep(0.2)
        else:
            chunks = chunk_markdown_text(content, limit=2000)
            total = len(chunks)
            for i, chunk in enumerate(chunks):
                if total > 1:
                    chunk = f"{chunk}\n({i + 1}/{total})"
                if chat_type == QQBotChatType.C2C:
                    await self._api_client.send_c2c_message(
                        target, chunk, msg_type=0,
                        msg_id=reply_to_id if i == 0 else None,
                        message_reference=message_reference if i == total - 1 else None,
                    )
                elif chat_type == QQBotChatType.GUILD:
                    await self._api_client.send_channel_message(
                        target, chunk, msg_type=0,
                    )
                else:
                    await self._api_client.send_group_message(
                        target, chunk, msg_type=0,
                        msg_id=reply_to_id if i == 0 else None,
                        message_reference=message_reference if i == total - 1 else None,
                    )
                if i < total - 1:
                    await asyncio.sleep(0.2)

    async def send_c2c_stream(
        self,
        openid: str,
        stream_fn: Any | None = None,
    ) -> str | None:
        if not self._streaming_controller:
            return None
        return await self._streaming_controller.start_and_stream(openid, stream_fn)

    async def send_media(
        self,
        target_id: str,
        media_url: str,
        media_type: str,
        reply_to_id: str | None = None,
        account_id: str | None = None,
    ) -> None:
        from yuxi.channel.extensions.qqbot.outbound_media import QQBotOutboundMedia
        from yuxi.channel.extensions.qqbot.media import QQBotMedia

        media_manager = QQBotMedia(
            api_client=self._api_client,
            url_direct_upload=self._account.url_direct_upload,
        )
        outbound_media = QQBotOutboundMedia(
            api_client=self._api_client,
            media_manager=media_manager,
        )
        send_map = {
            "image": outbound_media.send_photo,
            "voice": outbound_media.send_voice,
            "video": outbound_media.send_video,
            "file": outbound_media.send_document,
        }
        send_fn = send_map.get(media_type) or send_map.get("image")
        if send_fn is None:
            return
        await send_fn(target_id, media_url, reply_to_id=reply_to_id)

    async def send_input_notify(self, openid: str) -> None:
        await self._api_client.send_input_notify(openid)

    async def send_keyboard_message(
        self,
        target_id: str,
        content: str,
        keyboard: dict,
        reply_to_id: str | None = None,
    ) -> None:
        chat_type, target = self._parse_target(target_id)
        formatted = md_to_qq_markdown(content)
        if chat_type == QQBotChatType.C2C:
            await self._api_client.send_c2c_markdown(target, formatted, keyboard=keyboard, msg_id=reply_to_id)
        else:
            await self._api_client.send_group_markdown(target, formatted, keyboard=keyboard, msg_id=reply_to_id)

    async def unsend(self, target_id: str, message_id: str) -> bool:
        chat_type, target = self._parse_target(target_id)
        try:
            if chat_type == QQBotChatType.C2C:
                await self._api_client.delete_c2c_message(target, message_id)
            elif chat_type == QQBotChatType.GUILD:
                await self._api_client.delete_channel_message(target, message_id)
            else:
                await self._api_client.delete_group_message(target, message_id)
            return True
        except Exception:
            logger.exception("unsend failed: target=%s, msg_id=%s", target_id, message_id)
            return False

    async def send_ark(self, target_id: str, ark: dict) -> None:
        chat_type, target = self._parse_target(target_id)
        if chat_type == QQBotChatType.C2C:
            await self._api_client.send_c2c_ark(target, ark)
        else:
            await self._api_client.send_group_ark(target, ark)

    def _parse_target(self, target_id: str) -> tuple[QQBotChatType, str]:
        parts = target_id.split(":", 2)
        if len(parts) >= 3:
            prefix = parts[1]
            actual_id = parts[2]
            if prefix in ("c2c", "dm"):
                return (QQBotChatType.C2C, actual_id)
            elif prefix == "group":
                return (QQBotChatType.GROUP, actual_id)
            elif prefix == "channel":
                return (QQBotChatType.GUILD, actual_id)
        return (QQBotChatType.C2C, target_id)