ForcePilot/backend/package/yuxi/channels/adapters/wechat/wecom/send.py

from __future__ import annotations

from typing import Any

import httpx

from yuxi.channels.exceptions import ChannelRateLimitError
from yuxi.channels.models import DeliveryResult
from yuxi.channels.adapters.wechat.format import truncate_text
from yuxi.channels.adapters.wechat.errors import is_token_expired, parse_wecom_error
from yuxi.channels.adapters.wechat.retry import retry_with_backoff

from .client import WeComClient


async def send_wecom_message(
    client: WeComClient,
    http_client: httpx.AsyncClient,
    payload: dict[str, Any],
) -> DeliveryResult:
    token = await client.get_access_token()
    api_url = f"https://qyapi.weixin.qq.com/cgi-bin/message/send?access_token={token}"

    async def _post():
        return await http_client.post(api_url, json=payload)

    try:
        resp = await retry_with_backoff(_post)
        data = resp.json()

        if data.get("errcode") == 0:
            return DeliveryResult(success=True, message_id=data.get("msgid"))
        elif is_token_expired(data.get("errcode", 0)):
            client.invalidate_token()
            token = await client.get_access_token()
            api_url = f"https://qyapi.weixin.qq.com/cgi-bin/message/send?access_token={token}"

            async def _retry_post():
                return await http_client.post(api_url, json=payload)

            resp = await retry_with_backoff(_retry_post)
            data = resp.json()
            if data.get("errcode") == 0:
                return DeliveryResult(success=True, message_id=data.get("msgid"))
            _, err_detail = parse_wecom_error(data)
            return DeliveryResult(success=False, error=err_detail)
        elif data.get("errcode") in (45009, 45011):
            raise ChannelRateLimitError(retry_after_ms=60000)
        else:
            _, err_detail = parse_wecom_error(data)
            return DeliveryResult(success=False, error=err_detail)
    except httpx.HTTPError as e:
        return DeliveryResult(success=False, error=str(e))


def build_wecom_text_payload(
    agent_id: str,
    to_user: str,
    content: str,
    chat_type: str = "direct",
    reply_to_msg_id: str | None = None,
    reply_to_user: str | None = None,
    safe: int = 0,
) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "msgtype": "text",
        "agentid": agent_id,
        "text": {"content": truncate_text(content, 2048)},
        "safe": safe,
    }
    payload["touser"] = to_user
    if reply_to_msg_id:
        payload["_reply_to_msg_id"] = reply_to_msg_id
        if reply_to_user:
            payload["_reply_to_user"] = reply_to_user
            quoted_prefix = f"「回复 @{reply_to_user}」\n"
            if len(quoted_prefix + content) <= 2048:
                payload["text"]["content"] = quoted_prefix + content
    return payload


def build_wecom_image_payload(agent_id: str, to_user: str, media_id: str) -> dict[str, Any]:
    return {
        "touser": to_user,
        "msgtype": "image",
        "agentid": agent_id,
        "image": {"media_id": media_id},
    }


def build_wecom_file_payload(agent_id: str, to_user: str, media_id: str) -> dict[str, Any]:
    return {
        "touser": to_user,
        "msgtype": "file",
        "agentid": agent_id,
        "file": {"media_id": media_id},
    }


def build_wecom_voice_payload(agent_id: str, to_user: str, media_id: str) -> dict[str, Any]:
    return {
        "touser": to_user,
        "msgtype": "voice",
        "agentid": agent_id,
        "voice": {"media_id": media_id},
    }


def build_wecom_video_payload(
    agent_id: str, to_user: str, media_id: str, title: str = "", description: str = ""
) -> dict[str, Any]:
    return {
        "touser": to_user,
        "msgtype": "video",
        "agentid": agent_id,
        "video": {"media_id": media_id, "title": title, "description": description},
    }


def build_wecom_news_payload(agent_id: str, to_user: str, articles: list[dict[str, Any]]) -> dict[str, Any]:
    return {
        "touser": to_user,
        "msgtype": "news",
        "agentid": agent_id,
        "news": {"articles": articles},
    }


def build_wecom_miniprogram_payload(
    agent_id: str,
    to_user: str,
    title: str,
    appid: str,
    pagepath: str,
    thumb_media_id: str,
) -> dict[str, Any]:
    return {
        "touser": to_user,
        "msgtype": "miniprogram_notice",
        "agentid": agent_id,
        "miniprogram_notice": {
            "appid": appid,
            "title": title,
            "page": pagepath,
            "emphasis_first_item": False,
            "content_item": [{"key": "详情", "value": title}],
        },
    }


async def send_wecom_voice(
    client: WeComClient,
    http_client: httpx.AsyncClient,
    agent_id: str,
    to_user: str,
    voice_data: bytes,
) -> DeliveryResult:
    try:
        media_id = await client.upload_media(voice_data, "voice.amr", "voice")
    except Exception as e:
        return DeliveryResult(success=False, error=f"WeCom voice upload failed: {e}")

    payload = build_wecom_voice_payload(agent_id, to_user, media_id)
    return await send_wecom_message(client, http_client, payload)


async def send_wecom_video(
    client: WeComClient,
    http_client: httpx.AsyncClient,
    agent_id: str,
    to_user: str,
    video_data: bytes,
    title: str = "",
    description: str = "",
) -> DeliveryResult:
    try:
        media_id = await client.upload_media(video_data, "video.mp4", "video")
    except Exception as e:
        return DeliveryResult(success=False, error=f"WeCom video upload failed: {e}")

    payload = build_wecom_video_payload(agent_id, to_user, media_id, title, description)
    return await send_wecom_message(client, http_client, payload)