ForcePilot/backend/package/yuxi/channels/adapters/qqbot/format.py

from __future__ import annotations

from yuxi.channels.models import ChannelResponse

from .constants import DM_CHAT_PREFIX, GROUP_CHAT_PREFIX


def build_text_payload(response: ChannelResponse) -> dict:
    content = response.content[:2000]
    chat_id = response.identity.channel_chat_id

    payload: dict = {"content": content}

    payload["msg_type"] = 0

    if chat_id.startswith(GROUP_CHAT_PREFIX):
        payload["group_openid"] = chat_id.replace(GROUP_CHAT_PREFIX, "")
    elif not chat_id.startswith(DM_CHAT_PREFIX):
        payload["channel_id"] = chat_id

    if response.reply_to_message_id:
        payload["msg_id"] = response.reply_to_message_id

    return payload


def build_markdown_payload(response: ChannelResponse, template_id: str | None = None) -> dict:
    md_template_id = response.metadata.get("markdown_template_id") or template_id

    if md_template_id:
        return {
            "msg_type": 2,
            "markdown": {
                "template_id": md_template_id,
                "params": [
                    {"key": "title", "values": [response.metadata.get("title", "")]},
                    {"key": "content", "values": [response.content[:4096]]},
                ],
            },
        }

    return {
        "msg_type": 2,
        "markdown": {
            "content": response.content[:4096],
        },
    }


def build_ark_payload(response: ChannelResponse) -> dict:
    ark_template_id = response.metadata.get("ark_template_id")
    ark_data = response.metadata.get("ark_data", {})

    return {
        "msg_type": 3,
        "ark": {
            "template_id": ark_template_id,
            "kv": [{"key": k, "value": v} for k, v in ark_data.items()],
        },
    }


def build_embed_payload(response: ChannelResponse) -> dict:
    embed_data = response.metadata.get("embed", {})

    return {
        "msg_type": 4,
        "embed": {
            "title": embed_data.get("title", ""),
            "description": response.content[:4096],
            "prompt": embed_data.get("prompt", response.content[:200]),
            "fields": embed_data.get("fields", []),
        },
    }


def build_media_generic_payload(response: ChannelResponse, file_id: str, msg_type: int = 7) -> dict:
    chat_id = response.identity.channel_chat_id
    payload: dict = {
        "msg_type": msg_type,
        "media": {"file_info": file_id},
    }

    if response.content:
        payload["content"] = response.content[:2000]

    if chat_id.startswith(GROUP_CHAT_PREFIX):
        payload["group_openid"] = chat_id.replace(GROUP_CHAT_PREFIX, "")
    elif not chat_id.startswith(DM_CHAT_PREFIX):
        payload["channel_id"] = chat_id

    return payload


def build_image_payload(response: ChannelResponse, file_id: str) -> dict:
    chat_id = response.identity.channel_chat_id
    payload: dict = {
        "msg_type": 1,
        "image": file_id,
        "content": response.content[:2000] if response.content else "",
    }

    if chat_id.startswith(GROUP_CHAT_PREFIX):
        payload["group_openid"] = chat_id.replace(GROUP_CHAT_PREFIX, "")
    elif not chat_id.startswith(DM_CHAT_PREFIX):
        payload["channel_id"] = chat_id

    return payload


def format_outbound(
    response: ChannelResponse,
    use_markdown: bool = False,
    markdown_template_id: str | None = None,
) -> dict:
    msg_type = response.metadata.get("qq_msg_type", "")

    if msg_type == "markdown" or (use_markdown and not msg_type):
        return build_markdown_payload(response, markdown_template_id)
    elif msg_type == "ark" and response.metadata.get("ark_template_id"):
        return build_ark_payload(response)
    elif msg_type == "embed":
        return build_embed_payload(response)
    elif msg_type == "image" and response.attachments:
        file_id = response.attachments[0].file_id or response.attachments[0].url or ""
        return build_image_payload(response, file_id)
    elif msg_type in ("voice", "video", "file") and response.attachments:
        file_id_list = [
            response.metadata.get("media_file_id", ""),
            response.attachments[0].file_id or "",
            response.attachments[0].url or "",
        ]
        file_id = next((fid for fid in file_id_list if fid), "")
        return build_media_generic_payload(response, file_id, msg_type=7)
    else:
        return build_text_payload(response)


class MarkdownChunker:
    MAX_CHARS = 5000
    CHUNK_OVERLAP = 200

    def __init__(self, max_chars: int = MAX_CHARS, chunk_overlap: int = CHUNK_OVERLAP):
        self._max_chars = max_chars
        self._chunk_overlap = chunk_overlap

    def chunk(self, text: str) -> list[str]:
        if len(text) <= self._max_chars:
            return [text]

        paragraphs = self._split_paragraphs(text)
        chunks: list[str] = []
        current_chunk: list[str] = []
        current_len = 0

        for para in paragraphs:
            para_len = len(para)
            if current_len + para_len <= self._max_chars:
                current_chunk.append(para)
                current_len += para_len
            else:
                if current_chunk:
                    chunks.append("".join(current_chunk))
                if para_len > self._max_chars:
                    sub_chunks = self._force_split(para)
                    if current_chunk:
                        for i, sc in enumerate(sub_chunks):
                            chunks.append(sc)
                    else:
                        chunks.extend(sub_chunks)
                    current_chunk = []
                    current_len = 0
                else:
                    current_chunk = [para]
                    current_len = para_len

        if current_chunk:
            chunks.append("".join(current_chunk))

        return chunks

    def _split_paragraphs(self, text: str) -> list[str]:
        sections: list[str] = []
        in_code_block = False
        current: list[str] = []
        lines = text.splitlines(keepends=True)

        for line in lines:
            stripped = line.strip()
            if stripped.startswith("```"):
                if current:
                    sections.append("".join(current))
                    current = []
                if in_code_block:
                    sections.append(line)
                    in_code_block = False
                else:
                    in_code_block = True
                    current.append(line)
                continue

            if in_code_block:
                current.append(line)
                if stripped.endswith("```"):
                    sections.append("".join(current))
                    current = []
                    in_code_block = False
                continue

            if not stripped:
                if current:
                    sections.append("".join(current))
                    current = []
                sections.append(line)
            elif (
                stripped.startswith(("#", "-", "*", ">", "|"))
                and current
                and not current[-1].strip().startswith(("#", "-", "*", ">", "|", "1.", "2.", "3."))
            ):
                if current:
                    sections.append("".join(current))
                    current = []
                current.append(line)
            else:
                current.append(line)

        if current:
            sections.append("".join(current))

        result: list[str] = []
        buffer: list[str] = []
        for s in sections:
            stripped = s.strip()
            if not stripped and buffer:
                result.append("".join(buffer))
                buffer = []
            buffer.append(s)

        if buffer:
            content = "".join(buffer)
            if content.strip():
                result.append(content)

        return result or [text]

    def _force_split(self, text: str) -> list[str]:
        chunks: list[str] = []
        for i in range(0, len(text), self._max_chars - self._chunk_overlap):
            chunks.append(text[i : i + self._max_chars])
        return chunks