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

from __future__ import annotations

import asyncio
import logging
import os
import time
from pathlib import Path

import httpx
from requests_oauthlib import OAuth1Session

logger = logging.getLogger(__name__)

MEDIA_CATEGORY_SIZES = {
    "dm_image": 5 * 1024 * 1024,
    "dm_video": 512 * 1024 * 1024,
    "dm_gif": 15 * 1024 * 1024,
}

MEDIA_CATEGORY_TYPES = {
    "dm_image": "image",
    "dm_video": "video",
    "dm_gif": "image",
}

CHUNK_SIZE = 5 * 1024 * 1024

X_UPLOAD_BASE = "https://upload.twitter.com"


class TwitterMedia:
    @staticmethod
    def _build_oauth(account: dict) -> OAuth1Session:
        return OAuth1Session(
            account["api_key"],
            client_secret=account["api_secret"],
            resource_owner_key=account["access_token"],
            resource_owner_secret=account["access_secret"],
        )

    @classmethod
    async def upload(
        cls, account: dict, media_url_or_path: str, media_category: str
    ) -> str | None:
        max_size = MEDIA_CATEGORY_SIZES.get(media_category, 5 * 1024 * 1024)
        media_type = MEDIA_CATEGORY_TYPES.get(media_category, "image")

        file_path = await cls._resolve_file(media_url_or_path)
        if not file_path:
            logger.error(
                "Twitter media: cannot resolve file from %s", media_url_or_path
            )
            return None

        file_size = os.path.getsize(file_path)
        if file_size > max_size:
            logger.error(
                "Twitter media: file too large (%d bytes, max %d for %s)",
                file_size,
                max_size,
                media_category,
            )
            return None

        oauth = cls._build_oauth(account)

        init_result = await asyncio.to_thread(
            cls._init_upload,
            oauth,
            file_size,
            media_category,
            media_type,
        )
        if not init_result:
            return None

        media_id = init_result["media_id_string"]

        with open(file_path, "rb") as f:
            segment_index = 0
            while True:
                chunk = f.read(CHUNK_SIZE)
                if not chunk:
                    break
                success = await asyncio.to_thread(
                    cls._append_chunk,
                    oauth,
                    media_id,
                    segment_index,
                    chunk,
                )
                if not success:
                    logger.error(
                        "Twitter media: APPEND failed at segment %d", segment_index
                    )
                    return None
                segment_index += 1

        finalize_result = await asyncio.to_thread(
            cls._finalize_upload,
            oauth,
            media_id,
        )
        if not finalize_result:
            return None

        if media_category == "dm_video":
            media_id = await asyncio.to_thread(
                cls._wait_for_video_processing,
                oauth,
                media_id,
            )

        logger.info(
            "Twitter media: uploaded media_id=%s, size=%d bytes", media_id, file_size
        )
        return media_id

    @staticmethod
    def _init_upload(
        oauth: OAuth1Session, total_bytes: int, media_category: str, media_type: str
    ) -> dict | None:
        url = f"{X_UPLOAD_BASE}/1.1/media/upload.json"
        params = {
            "command": "INIT",
            "total_bytes": total_bytes,
            "media_type": media_type,
            "media_category": media_category,
        }
        try:
            resp = oauth.post(url, data=params, timeout=30)
            if resp.status_code in (200, 202):
                return resp.json()
            logger.error(
                "Twitter media INIT failed: HTTP %d %s",
                resp.status_code,
                resp.text[:200],
            )
        except Exception as e:
            logger.error("Twitter media INIT error: %s", e)
        return None

    @staticmethod
    def _append_chunk(
        oauth: OAuth1Session, media_id: str, segment_index: int, chunk: bytes
    ) -> bool:
        url = f"{X_UPLOAD_BASE}/1.1/media/upload.json"
        params = {
            "command": "APPEND",
            "media_id": media_id,
            "segment_index": segment_index,
        }
        try:
            resp = oauth.post(url, data=params, files={"media": chunk}, timeout=60)
            return resp.status_code in (200, 204)
        except Exception as e:
            logger.error(
                "Twitter media APPEND error (segment %d): %s", segment_index, e
            )
            return False

    @staticmethod
    def _finalize_upload(oauth: OAuth1Session, media_id: str) -> dict | None:
        url = f"{X_UPLOAD_BASE}/1.1/media/upload.json"
        params = {
            "command": "FINALIZE",
            "media_id": media_id,
        }
        try:
            resp = oauth.post(url, data=params, timeout=30)
            if resp.status_code in (200, 202):
                return resp.json()
            logger.error(
                "Twitter media FINALIZE failed: HTTP %d %s",
                resp.status_code,
                resp.text[:200],
            )
        except Exception as e:
            logger.error("Twitter media FINALIZE error: %s", e)
        return None

    @staticmethod
    def _wait_for_video_processing(
        oauth: OAuth1Session, media_id: str, max_wait_sec: int = 300
    ) -> str | None:
        url = f"{X_UPLOAD_BASE}/1.1/media/upload.json"
        params = {"command": "STATUS", "media_id": media_id}
        start = time.monotonic()
        while time.monotonic() - start < max_wait_sec:
            try:
                resp = oauth.get(url, params=params, timeout=15)
                if resp.status_code != 200:
                    break
                data = resp.json()
                processing_info = data.get("processing_info", {})
                state = processing_info.get("state", "succeeded")
                if state == "succeeded":
                    return media_id
                if state == "failed":
                    logger.error(
                        "Twitter media: video processing failed for %s", media_id
                    )
                    return None
                check_after = processing_info.get("check_after_secs", 5)
                time.sleep(check_after)
            except Exception as e:
                logger.error("Twitter media STATUS error: %s", e)
                return None
        logger.error("Twitter media: video processing timeout for %s", media_id)
        return None

    @staticmethod
    async def _resolve_file(media_url_or_path: str) -> str | None:
        if media_url_or_path.startswith(("http://", "https://")):
            try:
                async with httpx.AsyncClient(timeout=httpx.Timeout(30.0)) as client:
                    resp = await client.get(media_url_or_path)
                    if resp.status_code == 200:
                        import tempfile

                        suffix = ".jpg"
                        content_type = resp.headers.get("content-type", "")
                        if "png" in content_type:
                            suffix = ".png"
                        elif "gif" in content_type:
                            suffix = ".gif"
                        elif "mp4" in content_type:
                            suffix = ".mp4"
                        elif "webp" in content_type:
                            suffix = ".webp"
                        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
                        tmp.write(resp.content)
                        tmp.close()
                        return tmp.name
            except Exception as e:
                logger.error("Twitter media download error: %s", e)
                return None
        path = Path(media_url_or_path)
        if path.is_file():
            return str(path)
        return None