ForcePilot/backend/package/yuxi/channel/extensions/wechat-mp/media.py

import asyncio
import logging
import time

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"
MATERIAL_ADD_URL = "https://api.weixin.qq.com/cgi-bin/material/add_material"
MATERIAL_DEL_URL = "https://api.weixin.qq.com/cgi-bin/material/del_material"
CLEAR_QUOTA_URL = "https://api.weixin.qq.com/cgi-bin/clear_quota"
MATERIAL_BATCHGET_URL = "https://api.weixin.qq.com/cgi-bin/material/batchget_material"
MATERIAL_COUNT_URL = "https://api.weixin.qq.com/cgi-bin/material/get_materialcount"
MATERIAL_ADD_NEWS_URL = "https://api.weixin.qq.com/cgi-bin/material/add_news"

ALLOWED_IMAGE_TYPES = {"image/jpeg", "image/jpg", "image/png"}
ALLOWED_VOICE_TYPES = {"audio/amr", "audio/mp3", "audio/speex"}
ALLOWED_VIDEO_TYPES = {"video/mp4"}
ALLOWED_THUMB_TYPES = {"image/jpeg", "image/jpg"}

ALLOWED_MEDIA_TYPE_MAP = {
    "image": ALLOWED_IMAGE_TYPES,
    "voice": ALLOWED_VOICE_TYPES,
    "video": ALLOWED_VIDEO_TYPES,
    "thumb": ALLOWED_THUMB_TYPES,
}

MAX_IMAGE_BYTES = 10 * 1024 * 1024
MAX_VOICE_BYTES = 2 * 1024 * 1024
MAX_VIDEO_BYTES = 10 * 1024 * 1024
MAX_THUMB_BYTES = 64 * 1024

MAX_BYTES_MAP = {
    "image": MAX_IMAGE_BYTES,
    "voice": MAX_VOICE_BYTES,
    "video": MAX_VIDEO_BYTES,
    "thumb": MAX_THUMB_BYTES,
}


class WeChatMedia:

    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
        self._last_clear_quota = 0.0

    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(
                    "WeChat 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("WeChat 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(
                        "WeChat 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("WeChat media download exception")
            return {"success": False, "error": str(e)}

    async def upload_permanent(self, file_data: bytes, filename: str, media_type: str) -> dict:
        token = self._token_provider()
        if not token:
            return {"success": False, "error": "no access_token"}

        url = f"{MATERIAL_ADD_URL}?access_token={token}&type={media_type}"
        files = {"media": (filename, file_data, "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"], "url": data.get("url", "")}
            if "errcode" in data:
                logger.warning(
                    "Permanent material 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("Permanent material upload exception")
            return {"success": False, "error": str(e)}

    async def delete_permanent(self, media_id: str) -> bool:
        token = self._token_provider()
        if not token:
            return False

        url = f"{MATERIAL_DEL_URL}?access_token={token}"
        try:
            client = await self._client()
            resp = await client.post(url, json={"media_id": media_id})
            data = resp.json()
            return data.get("errcode") == 0
        except Exception:
            logger.exception("Permanent material delete exception")
            return False

    async def delete_after_delay(self, media_id: str, delay: float = 10.0):
        await asyncio.sleep(delay)
        await self.delete_permanent(media_id)

    async def clear_quota(self) -> bool:
        now = time.time()
        if now - self._last_clear_quota < 60:
            return False
        self._last_clear_quota = now

        if not self._app_id or not self._app_secret:
            return False

        payload = {"appid": self._app_id, "appsecret": self._app_secret}
        try:
            client = await self._client()
            resp = await client.post(CLEAR_QUOTA_URL, json=payload)
            data = resp.json()
            ok = data.get("errcode") == 0
            if not ok:
                logger.warning("Clear quota failed: errcode=%s errmsg=%s", data.get("errcode"), data.get("errmsg"))
            return ok
        except Exception:
            logger.exception("Clear quota exception")
            return False

    async def list_materials(self, media_type: str, offset: int = 0, count: int = 20) -> dict:
        token = self._token_provider()
        if not token:
            return {"success": False, "error": "no access_token"}

        payload = {"type": media_type, "offset": offset, "count": min(count, 20)}
        try:
            client = await self._client()
            resp = await client.post(f"{MATERIAL_BATCHGET_URL}?access_token={token}", json=payload)
            data = resp.json()
            if data.get("errcode") == 0:
                return {
                    "success": True,
                    "total_count": data.get("total_count", 0),
                    "item_count": data.get("item_count", 0),
                    "item": data.get("item", []),
                }
            return {"success": False, "error": data.get("errmsg", "list failed"), "errcode": data.get("errcode")}
        except Exception as e:
            logger.exception("List materials exception")
            return {"success": False, "error": str(e)}

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

        try:
            client = await self._client()
            resp = await client.get(f"{MATERIAL_COUNT_URL}?access_token={token}")
            data = resp.json()
            if data.get("errcode") == 0:
                return {
                    "success": True,
                    "voice_count": data.get("voice_count", 0),
                    "video_count": data.get("video_count", 0),
                    "image_count": data.get("image_count", 0),
                    "news_count": data.get("news_count", 0),
                }
            return {"success": False, "error": data.get("errmsg", "count failed"), "errcode": data.get("errcode")}
        except Exception as e:
            logger.exception("Get material count exception")
            return {"success": False, "error": str(e)}

    async def upload_news(self, articles: list[dict]) -> dict:
        token = self._token_provider()
        if not token:
            return {"success": False, "error": "no access_token"}

        payload = {"articles": articles}
        try:
            client = await self._client()
            resp = await client.post(f"{MATERIAL_ADD_NEWS_URL}?access_token={token}", json=payload)
            data = resp.json()
            if "media_id" in data:
                return {"success": True, "media_id": data["media_id"]}
            return {"success": False, "error": data.get("errmsg", "upload news failed"), "errcode": data.get("errcode")}
        except Exception as e:
            logger.exception("Upload news exception")
            return {"success": False, "error": str(e)}

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