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

import logging

import httpx

logger = logging.getLogger(__name__)

UPLOAD_URL = "https://api.weixin.qq.com/cgi-bin/media/upload"
DOWNLOAD_URL = "https://api.weixin.qq.com/cgi-bin/media/get"

ALLOWED_MEDIA_TYPE_MAP = {
    "image": {"image/jpeg", "image/jpg", "image/png"},
    "voice": {"audio/amr", "audio/mp3", "audio/speex"},
    "video": {"video/mp4"},
    "thumb": {"image/jpeg", "image/jpg"},
}

MAX_BYTES_MAP = {
    "image": 10 * 1024 * 1024,
    "voice": 2 * 1024 * 1024,
    "video": 10 * 1024 * 1024,
    "thumb": 64 * 1024,
}


class MiniProgramMedia:

    def __init__(self, access_token_provider, app_id: str = "", app_secret: str = ""):
        self._token_provider = access_token_provider
        self._http: httpx.AsyncClient | None = None
        self._app_id = app_id
        self._app_secret = app_secret

    async def _client(self) -> httpx.AsyncClient:
        if self._http is None:
            self._http = httpx.AsyncClient(timeout=30.0)
        return self._http

    async def upload(
        self,
        file_data: bytes,
        filename: str,
        media_type: str,
        content_type: str | None = None,
    ) -> dict:
        if media_type not in ALLOWED_MEDIA_TYPE_MAP:
            return {"success": False, "error": f"unknown media_type: {media_type}"}

        max_bytes = MAX_BYTES_MAP.get(media_type, 0)
        if len(file_data) > max_bytes:
            return {"success": False, "error": f"file too large: {len(file_data)} bytes (max {max_bytes})"}

        token = self._token_provider()
        if not token:
            return {"success": False, "error": "no access_token"}

        url = f"{UPLOAD_URL}?access_token={token}&type={media_type}"
        files = {"media": (filename, file_data, content_type or "application/octet-stream")}

        try:
            client = await self._client()
            resp = await client.post(url, files=files)
            data = resp.json()

            if "media_id" in data:
                return {"success": True, "media_id": data["media_id"], "type": data.get("type", media_type)}
            if "errcode" in data:
                logger.warning(
                    "MiniProgram media upload failed: errcode=%s errmsg=%s",
                    data.get("errcode"),
                    data.get("errmsg"),
                )
                return {"success": False, "error": data.get("errmsg", "upload failed"), "errcode": data["errcode"]}
            return {"success": False, "error": "unknown response"}
        except Exception as e:
            logger.exception("MiniProgram media upload exception")
            return {"success": False, "error": str(e)}

    async def download(self, media_id: str) -> dict:
        token = self._token_provider()
        if not token:
            return {"success": False, "error": "no access_token"}

        url = f"{DOWNLOAD_URL}?access_token={token}&media_id={media_id}"

        try:
            client = await self._client()
            resp = await client.get(url)

            content_type = resp.headers.get("content-type", "")
            if "application/json" in content_type or resp.text.startswith("{"):
                data = resp.json()
                if "errcode" in data and data["errcode"] != 0:
                    logger.warning(
                        "MiniProgram media download failed: errcode=%s errmsg=%s",
                        data.get("errcode"),
                        data.get("errmsg"),
                    )
                    return {
                        "success": False,
                        "error": data.get("errmsg", "download failed"),
                        "errcode": data["errcode"],
                    }
                return {"success": False, "error": "unexpected json response"}

            return {
                "success": True,
                "data": resp.content,
                "content_type": content_type,
            }

        except Exception as e:
            logger.exception("MiniProgram media download exception")
            return {"success": False, "error": str(e)}

    async def download_image_base64(self, media_id: str) -> str:
        result = await self.download(media_id)
        if result.get("success") and result.get("data"):
            import base64

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

    async def close(self):
        if self._http:
            await self._http.aclose()
            self._http = None