ForcePilot/backend/package/yuxi/channel/extensions/wechat_miniprogram/message.py

import xml.etree.ElementTree as ET
from datetime import datetime, UTC

from yuxi.channel.extensions.wechat_mp.types import InboundWeChatMessage
from yuxi.channel.message.models import MessageType, PeerInfo, UnifiedMessage
from yuxi.channel.routing.models import PeerKind


def parse_xml_to_message(xml_text: str) -> InboundWeChatMessage:
    root = ET.fromstring(xml_text)

    def _text(tag: str) -> str:
        el = root.find(tag)
        return el.text or "" if el is not None else ""

    def _int(tag: str) -> int:
        try:
            return int(_text(tag))
        except (ValueError, TypeError):
            return 0

    def _float(tag: str) -> float:
        try:
            return float(_text(tag))
        except (ValueError, TypeError):
            return 0.0

    msg_type = _text("MsgType")
    event = _text("Event") if msg_type == "event" else ""

    msg_id = _text("MsgId")
    if not msg_id:
        msg_id = f"{_text('FromUserName')}_{_text('CreateTime')}_{msg_type}_{event}"

    return InboundWeChatMessage(
        msg_id=msg_id,
        msg_type=msg_type,
        from_user=_text("FromUserName"),
        to_user=_text("ToUserName"),
        create_time=_int("CreateTime"),
        content=_text("Content"),
        pic_url=_text("PicUrl"),
        media_id=_text("MediaId"),
        media_id_16k=_text("MediaId16K"),
        thumb_media_id=_text("ThumbMediaId"),
        media_format=_text("Format"),
        recognition=_text("Recognition"),
        location_x=_float("Location_X"),
        location_y=_float("Location_Y"),
        label=_text("Label"),
        title=_text("Title"),
        description=_text("Description"),
        url=_text("Url"),
        event=event,
        event_key=_text("EventKey"),
        app_id=_text("AppId"),
        page_path=_text("PagePath"),
        thumb_url=_text("ThumbUrl"),
        raw_xml=_dict_from_element(root),
    )


def resolve_voice(msg: InboundWeChatMessage) -> dict:
    if msg.recognition and msg.recognition.strip():
        return {"mode": "recognition", "text": msg.recognition, "ctype": "text"}
    if msg.media_id:
        return {"mode": "download", "media_id": msg.media_id, "format": msg.media_format, "ctype": "voice"}
    return {"mode": "unknown", "ctype": "text"}


def extract_content(msg: InboundWeChatMessage) -> str:
    match msg.msg_type:
        case "text":
            return msg.content
        case "image":
            if msg.pic_url:
                return f"[图片: {msg.pic_url}]"
            if msg.media_id:
                return f"[图片: {msg.media_id}]"
            return "[图片]"
        case "voice":
            voice_info = resolve_voice(msg)
            if voice_info["mode"] == "recognition":
                return voice_info["text"]
            if voice_info["mode"] == "download":
                return f"[语音: {voice_info['media_id']}]"
            return "[语音]"
        case "video" | "shortvideo":
            if msg.media_id:
                return f"[视频: {msg.media_id}]"
            return "[视频]"
        case "location":
            return f"{msg.label}\n({msg.location_y}, {msg.location_x})"
        case "link":
            return f"{msg.title}\n{msg.description}\n{msg.url}"
        case "miniprogrampage":
            parts = []
            if msg.title:
                parts.append(f"小程序卡片: {msg.title}")
            if msg.app_id:
                parts.append(f"AppId: {msg.app_id}")
            if msg.page_path:
                parts.append(f"页面路径: {msg.page_path}")
            return "\n".join(parts) if parts else "[小程序卡片]"
        case _:
            return msg.content or ""


async def download_media(media_id: str, gateway) -> bytes | None:
    from yuxi.channel.extensions.wechat_mp.media import WeChatMedia

    token = gateway.access_token if gateway else None
    if not token:
        return None
    wm = WeChatMedia(lambda: token)
    result = await wm.download(media_id)
    if result.get("success"):
        return result["data"]
    return None


async def download_image_content(msg: InboundWeChatMessage, gateway) -> str:
    if not msg.media_id and not msg.pic_url:
        return ""
    if msg.pic_url:
        return msg.pic_url
    data = await download_media(msg.media_id, gateway)
    if data:
        import base64

        return f"data:image/jpeg;base64,{base64.b64encode(data).decode()}"
    return ""


def build_unified_message(raw_msg: InboundWeChatMessage, content: str, account_id: str = "default") -> UnifiedMessage:
    msg_type = MessageType.TEXT
    if raw_msg.msg_type == "image":
        msg_type = MessageType.IMAGE
    elif raw_msg.msg_type == "voice":
        msg_type = MessageType.VOICE

    return UnifiedMessage(
        msg_id=raw_msg.msg_id,
        channel_type="wechat-miniprogram",
        account_id=account_id,
        content=content,
        message_type=msg_type,
        sender=PeerInfo(
            id=raw_msg.from_user,
            kind=PeerKind.DIRECT,
            display_name=raw_msg.from_user,
        ),
        timestamp=datetime.fromtimestamp(raw_msg.create_time, tz=UTC) if raw_msg.create_time else None,
        raw_payload=raw_msg.raw_xml,
        body_for_agent=content,
        metadata={
            "FromUserName": raw_msg.from_user,
            "ToUserName": raw_msg.to_user,
            "MsgType": raw_msg.msg_type,
            "Event": raw_msg.event,
            "EventKey": raw_msg.event_key,
            "MediaId": raw_msg.media_id,
            "PicUrl": raw_msg.pic_url,
        },
    )


def _dict_from_element(element: ET.Element) -> dict:
    result = {}
    for child in element:
        if len(child) > 0:
            result[child.tag] = _dict_from_element(child)
        else:
            result[child.tag] = child.text or ""
    return result