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

import asyncio
import logging
import time

import httpx

logger = logging.getLogger(__name__)

USER_INFO_URL = "https://api.weixin.qq.com/cgi-bin/user/info"
USER_INFO_BATCH_URL = "https://api.weixin.qq.com/cgi-bin/user/info/batchget"
USER_REMARK_URL = "https://api.weixin.qq.com/cgi-bin/user/info/updateremark"
USER_LIST_URL = "https://api.weixin.qq.com/cgi-bin/user/get"
TAGS_URL = "https://api.weixin.qq.com/cgi-bin/tags"
BLACKLIST_URL = "https://api.weixin.qq.com/cgi-bin/tags/members/getblacklist"
BLACKLIST_BATCH_URL = "https://api.weixin.qq.com/cgi-bin/tags/members/batchblacklist"


class WeChatUserInfoCache:
    TTL_SECONDS = 3600

    def __init__(self):
        self._cache: dict[str, dict] = {}
        self._timestamps: dict[str, float] = {}
        self._lock = asyncio.Lock()

    async def get(self, openid: str) -> dict | None:
        async with self._lock:
            ts = self._timestamps.get(openid, 0)
            if time.time() - ts > self.TTL_SECONDS:
                self._cache.pop(openid, None)
                self._timestamps.pop(openid, None)
                return None
            return self._cache.get(openid)

    async def set(self, openid: str, info: dict) -> None:
        async with self._lock:
            self._cache[openid] = info
            self._timestamps[openid] = time.time()

    async def clear(self, openid: str) -> None:
        async with self._lock:
            self._cache.pop(openid, None)
            self._timestamps.pop(openid, None)


_user_cache = WeChatUserInfoCache()


async def fetch_user_info(openid: str, gateway) -> dict | None:
    cached = await _user_cache.get(openid)
    if cached:
        return cached

    token = gateway.access_token if gateway else None
    if not token:
        return None

    url = f"{USER_INFO_URL}?access_token={token}&openid={openid}&lang=zh_CN"
    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.get(url)
            data = resp.json()
            if data.get("errcode") == 0:
                info = {
                    "nickname": data.get("nickname", ""),
                    "headimgurl": data.get("headimgurl", ""),
                    "subscribe": data.get("subscribe", 0),
                    "sex": data.get("sex", 0),
                    "country": data.get("country", ""),
                    "province": data.get("province", ""),
                    "city": data.get("city", ""),
                }
                await _user_cache.set(openid, info)
                return info
            logger.warning("Fetch user info failed: errcode=%s errmsg=%s", data.get("errcode"), data.get("errmsg"))
    except Exception:
        logger.exception("Fetch user info exception for openid=%s", openid)
    return None


async def batch_fetch_user_info(openids: list[str], gateway) -> list[dict]:
    token = gateway.access_token if gateway else None
    if not token or not openids:
        return []

    payload = {"user_list": [{"openid": oid, "lang": "zh_CN"} for oid in openids[:100]]}
    try:
        async with httpx.AsyncClient(timeout=15.0) as client:
            resp = await client.post(f"{USER_INFO_BATCH_URL}?access_token={token}", json=payload)
            data = resp.json()
            if data.get("errcode") == 0:
                return data.get("user_info_list", [])
            logger.warning(
                "Batch fetch user info failed: errcode=%s errmsg=%s",
                data.get("errcode"),
                data.get("errmsg"),
            )
    except Exception:
        logger.exception("Batch fetch user info exception")
    return []


async def update_remark(openid: str, remark: str, gateway) -> bool:
    token = gateway.access_token if gateway else None
    if not token:
        return False

    payload = {"openid": openid, "remark": remark}
    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.post(f"{USER_REMARK_URL}?access_token={token}", json=payload)
            data = resp.json()
            return data.get("errcode") == 0
    except Exception:
        logger.exception("Update remark exception for openid=%s", openid)
    return False


async def fetch_followers(next_openid: str = "", gateway=None) -> dict:
    token = gateway.access_token if gateway else None
    if not token:
        return {"success": False, "error": "no access_token"}

    url = f"{USER_LIST_URL}?access_token={token}"
    if next_openid:
        url += f"&next_openid={next_openid}"
    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.get(url)
            data = resp.json()
            if data.get("errcode") == 0:
                return {
                    "success": True,
                    "total": data.get("total", 0),
                    "count": data.get("count", 0),
                    "data": data.get("data", {}).get("openid", []),
                    "next_openid": data.get("next_openid", ""),
                }
            return {"success": False, "error": data.get("errmsg", "fetch failed"), "errcode": data.get("errcode")}
    except Exception as e:
        logger.exception("Fetch followers exception")
        return {"success": False, "error": str(e)}


async def list_tags(gateway) -> list[dict]:
    token = gateway.access_token if gateway else None
    if not token:
        return []

    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.get(f"{TAGS_URL}/get?access_token={token}")
            data = resp.json()
            if data.get("errcode") == 0:
                return data.get("tags", [])
            logger.warning("List tags failed: errcode=%s errmsg=%s", data.get("errcode"), data.get("errmsg"))
    except Exception:
        logger.exception("List tags exception")
    return []


async def create_tag(name: str, gateway) -> int | None:
    token = gateway.access_token if gateway else None
    if not token:
        return None

    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.post(f"{TAGS_URL}/create?access_token={token}", json={"tag": {"name": name}})
            data = resp.json()
            if data.get("errcode") == 0:
                return data.get("tag", {}).get("id")
            logger.warning("Create tag failed: errcode=%s errmsg=%s", data.get("errcode"), data.get("errmsg"))
    except Exception:
        logger.exception("Create tag exception")
    return None


async def update_tag(tag_id: int, name: str, gateway) -> bool:
    token = gateway.access_token if gateway else None
    if not token:
        return False

    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.post(
                f"{TAGS_URL}/update?access_token={token}", json={"tag": {"id": tag_id, "name": name}}
            )
            data = resp.json()
            return data.get("errcode") == 0
    except Exception:
        logger.exception("Update tag exception")
    return False


async def delete_tag(tag_id: int, gateway) -> bool:
    token = gateway.access_token if gateway else None
    if not token:
        return False

    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.post(f"{TAGS_URL}/delete?access_token={token}", json={"tag": {"id": tag_id}})
            data = resp.json()
            return data.get("errcode") == 0
    except Exception:
        logger.exception("Delete tag exception")
    return False


async def tag_users(openids: list[str], tag_id: int, gateway) -> bool:
    token = gateway.access_token if gateway else None
    if not token:
        return False

    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.post(
                f"{TAGS_URL}/members/batchtagging?access_token={token}",
                json={"openid_list": openids[:50], "tagid": tag_id},
            )
            data = resp.json()
            return data.get("errcode") == 0
    except Exception:
        logger.exception("Tag users exception")
    return False


async def untag_users(openids: list[str], tag_id: int, gateway) -> bool:
    token = gateway.access_token if gateway else None
    if not token:
        return False

    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.post(
                f"{TAGS_URL}/members/batchuntagging?access_token={token}",
                json={"openid_list": openids[:50], "tagid": tag_id},
            )
            data = resp.json()
            return data.get("errcode") == 0
    except Exception:
        logger.exception("Untag users exception")
    return False


async def get_blacklist(begin_openid: str = "", gateway=None) -> dict:
    token = gateway.access_token if gateway else None
    if not token:
        return {"success": False, "error": "no access_token"}

    payload = {"begin_openid": begin_openid} if begin_openid else {}
    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.post(f"{BLACKLIST_URL}?access_token={token}", json=payload)
            data = resp.json()
            if data.get("errcode") == 0:
                return {
                    "success": True,
                    "total": data.get("total", 0),
                    "count": data.get("count", 0),
                    "data": data.get("data", {}).get("openid", []),
                    "next_openid": data.get("next_openid", ""),
                }
            return {"success": False, "error": data.get("errmsg", "fetch failed"), "errcode": data.get("errcode")}
    except Exception as e:
        logger.exception("Get blacklist exception")
        return {"success": False, "error": str(e)}


async def batch_blacklist(openids: list[str], gateway, action: str = "blacklist") -> bool:
    token = gateway.access_token if gateway else None
    if not token:
        return False

    url = f"{BLACKLIST_BATCH_URL}?access_token={token}"
    if action == "unblacklist":
        url = url.replace("batchblacklist", "batchunblacklist")
    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.post(url, json={"openid_list": openids[:20]})
            data = resp.json()
            return data.get("errcode") == 0
    except Exception:
        logger.exception("Batch blacklist exception")
    return False