ForcePilot/backend/package/yuxi/channel/extensions/xiaohongshu/outbound.py

import asyncio
import logging

import httpx

from yuxi.channel.extensions.xiaohongshu.format import clean_for_xiaohongshu
from yuxi.channel.extensions.xiaohongshu.window import XiaohongshuWindowTracker

logger = logging.getLogger(__name__)

SEND_MSG_PATH = "/api/message/send"
UPLOAD_MEDIA_PATH = "/api/material/upload"
MAX_TEXT_LEN = 2000

_send_context: dict[str, dict] = {}
_window_tracker = XiaohongshuWindowTracker()


def set_send_context(open_id: str, conversation_id: str, server_message_id: str) -> None:
    _send_context[open_id] = {
        "conversation_id": conversation_id,
        "server_message_id": server_message_id,
    }


def get_window_tracker() -> XiaohongshuWindowTracker:
    return _window_tracker


class XiaohongshuOutbound:

    def __init__(self, gateway=None):
        self._gateway = gateway

    def _http(self) -> httpx.AsyncClient | None:
        if self._gateway and hasattr(self._gateway, "_http"):
            return self._gateway._http
        return None

    async def send_text(
        self,
        to_user_id: str,
        content: str,
        *,
        reply_to_id: str | None = None,
        thread_id: str | None = None,
        account_id: str | None = None,
    ) -> bool:
        if not content:
            return False

        if not _window_tracker.can_reply(to_user_id):
            return False

        content = clean_for_xiaohongshu(content)
        if len(content) > MAX_TEXT_LEN:
            content = content[:MAX_TEXT_LEN]

        result = await self._send_msg(to_user_id, {
            "msg_type": "text",
            "content": {"text": content},
        })

        if result:
            _window_tracker.record_send(to_user_id)

        return result

    async def send_image(
        self,
        to_user_id: str,
        media_id: str,
    ) -> bool:
        if not media_id:
            return False

        if not _window_tracker.can_reply(to_user_id):
            return False

        result = await self._send_msg(to_user_id, {
            "msg_type": "image",
            "content": {"image_url": media_id},
        })

        if result:
            _window_tracker.record_send(to_user_id)

        return result

    async def send_video(
        self,
        to_user_id: str,
        media_id: str,
    ) -> bool:
        if not media_id:
            return False

        if not _window_tracker.can_reply(to_user_id):
            return False

        result = await self._send_msg(to_user_id, {
            "msg_type": "video",
            "content": {"video_url": media_id},
        })

        if result:
            _window_tracker.record_send(to_user_id)

        return result

    async def send_card(
        self,
        to_user_id: str,
        card_type: str,
        card_data: dict | None = None,
    ) -> bool:
        if not card_type:
            return False

        if not _window_tracker.can_reply(to_user_id):
            return False

        result = await self._send_msg(to_user_id, {
            "msg_type": "card",
            "content": {
                "card_type": card_type,
                "card_data": card_data or {},
            },
        })

        if result:
            _window_tracker.record_send(to_user_id)

        return result

    async def _send_msg(self, to_user_id: str, content: dict) -> bool:
        token = self._gateway.access_token if self._gateway else None
        if not token:
            logger.error("No valid access_token for xiaohongshu send")
            return False

        ctx = _send_context.pop(to_user_id, {})
        conversation_id = ctx.get("conversation_id", "")
        server_message_id = ctx.get("server_message_id", "")

        payload = {
            "access_token": token,
            "open_id": to_user_id,
            **content,
        }

        if conversation_id:
            payload["conversation_id"] = conversation_id
        if server_message_id:
            payload["msg_id"] = server_message_id

        headers = {
            "Content-Type": "application/json",
        }

        client = self._http()
        if client is None:
            logger.error("No HTTP client available for xiaohongshu send")
            return False

        for attempt in range(3):
            try:
                resp = await client.post(SEND_MSG_PATH, json=payload, headers=headers)
                data = resp.json()
                if data.get("code") == 0:
                    return True

                logger.warning(
                    "Xiaohongshu send failed (attempt %d): code=%s, msg=%s",
                    attempt + 1,
                    data.get("code"),
                    data.get("msg", ""),
                )

            except Exception as e:
                logger.warning("Xiaohongshu send attempt %d failed: %s", attempt + 1, e)

            await asyncio.sleep(attempt + 1)

        return False

    async def upload_image(self, image_url: str) -> str | None:
        token = self._gateway.access_token if self._gateway else None
        if not token:
            return None

        client = self._http()
        if client is None:
            logger.error("No HTTP client available for xiaohongshu upload")
            return None

        headers = {
            "Content-Type": "application/json",
        }

        for attempt in range(3):
            try:
                resp = await client.post(
                    UPLOAD_MEDIA_PATH,
                    json={"access_token": token, "image_url": image_url},
                    headers=headers,
                )
                data = resp.json()
                if data.get("code") == 0:
                    return data.get("data", {}).get("media_id")
                logger.warning("Xiaohongshu image upload failed (attempt %d): %s", attempt + 1, data)
            except Exception:
                logger.exception("Xiaohongshu image upload exception (attempt %d)", attempt + 1)

            await asyncio.sleep(attempt + 1)

        return None