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

from __future__ import annotations

import re

from yuxi.channels.models import MessageType

TRUNCATION_MARKER = "\n...(内容过长已截断)"

_MARKDOWN_BOUNDARY_PATTERN = re.compile(
    r'(\n\n|\n(?=#{1,6}\s)|\n(?=>\s)|\n(?=-\s)|\n(?=\d+\.\s))'
)


def truncate_text(text: str, limit: int) -> str:
    if len(text) <= limit:
        return text
    cut = max(0, limit - len(TRUNCATION_MARKER))
    return text[:cut] + TRUNCATION_MARKER


def truncate_markdown(text: str, limit: int) -> str:
    if len(text) <= limit:
        return text
    cut = max(0, limit - len(TRUNCATION_MARKER))
    candidates = list(_MARKDOWN_BOUNDARY_PATTERN.finditer(text[: cut + 1]))
    if candidates:
        boundary = candidates[-1].start()
        if boundary > cut * 0.7:
            return text[:boundary] + TRUNCATION_MARKER
    return text[:cut] + TRUNCATION_MARKER


def map_wecom_msg_type(msg_type: str) -> MessageType:
    mapping = {
        "text": MessageType.TEXT,
        "image": MessageType.IMAGE,
        "voice": MessageType.AUDIO,
        "video": MessageType.VIDEO,
        "file": MessageType.FILE,
        "location": MessageType.LOCATION,
    }
    return mapping.get(msg_type, MessageType.TEXT)


def map_mp_msg_type(msg_type: str) -> MessageType:
    mapping = {
        "text": MessageType.TEXT,
        "image": MessageType.IMAGE,
        "voice": MessageType.AUDIO,
        "video": MessageType.VIDEO,
        "location": MessageType.LOCATION,
    }
    return mapping.get(msg_type, MessageType.TEXT)


def map_bridge_msg_type(msg_type: int) -> MessageType:
    mapping = {
        1: MessageType.TEXT,
        3: MessageType.IMAGE,
        34: MessageType.AUDIO,
        43: MessageType.VIDEO,
        48: MessageType.LOCATION,
        49: MessageType.FILE,
    }
    return mapping.get(msg_type, MessageType.TEXT)