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

from __future__ import annotations

import logging
import os
import tempfile

logger = logging.getLogger(__name__)

TELEGRAM_API_BASE = "https://api.telegram.org"

MEDIA_TYPE_DETECT = {
    "image/jpeg": "photo",
    "image/png": "photo",
    "image/gif": "animation",
    "image/webp": "sticker",
    "video/mp4": "video",
    "audio/mpeg": "audio",
    "audio/ogg": "voice",
    "audio/wav": "voice",
}

TELEGRAM_MEDIA_LIMIT_MB = 50


class TelegramMedia:
    _file_service: object | None = None

    @classmethod
    def set_file_service(cls, file_service: object) -> None:
        cls._file_service = file_service

    @classmethod
    def classify_file(cls, filename: str) -> str:
        if cls._file_service is not None:
            return cls._file_service.classify(filename)
        return _classify_by_extension(filename)

    @staticmethod
    def detect_media_type(mime_type: str) -> str:
        return MEDIA_TYPE_DETECT.get(mime_type, "document")

    @staticmethod
    def is_within_size_limit(file_size_bytes: int, limit_mb: int = TELEGRAM_MEDIA_LIMIT_MB) -> bool:
        return file_size_bytes <= limit_mb * 1024 * 1024

    @classmethod
    async def download_file(cls, token: str, file_id: str, dest_dir: str | None = None) -> str | None:
        import httpx

        file_path_info = await cls._get_file_path(token, file_id)
        if not file_path_info:
            return None

        original_name = file_path_info.get("file_path", "").replace("/", "_")

        try:
            async with httpx.AsyncClient(timeout=httpx.Timeout(60.0)) as client:
                file_url = f"{TELEGRAM_API_BASE}/file/bot{token}/{file_path_info['file_path']}"
                resp = await client.get(file_url)
                if resp.status_code == 200:
                    if cls._file_service is not None:
                        safe_name = cls._file_service.generate_safe_name(original_name)
                        saved = await cls._file_service.save_upload(resp.content, safe_name)
                        return str(saved)

                    dest = dest_dir or tempfile.gettempdir()
                    local_path = os.path.join(dest, original_name)
                    os.makedirs(dest, exist_ok=True)
                    with open(local_path, "wb") as f:
                        f.write(resp.content)
                    return local_path
        except Exception:
            logger.exception("Failed to download Telegram file: %s", file_id)
        return None

    @staticmethod
    async def download_to_memory(token: str, file_id: str) -> bytes | None:
        import httpx

        file_path_info = await TelegramMedia._get_file_path(token, file_id)
        if not file_path_info:
            return None

        try:
            async with httpx.AsyncClient(timeout=httpx.Timeout(60.0)) as client:
                file_url = f"{TELEGRAM_API_BASE}/file/bot{token}/{file_path_info['file_path']}"
                resp = await client.get(file_url)
                if resp.status_code == 200:
                    return resp.content
        except Exception:
            logger.exception("Failed to download Telegram file to memory: %s", file_id)
        return None

    @staticmethod
    async def send_file(
        outbound, token: str, target_id: str, local_path: str, caption: str = "",
        thread_id: str | None = None,
    ) -> dict | None:
        import httpx

        filename = os.path.basename(local_path)
        mime_type = _guess_mime(filename)

        async with httpx.AsyncClient(timeout=httpx.Timeout(120.0)) as client:
            try:
                with open(local_path, "rb") as f:
                    method = "sendDocument"
                    data: dict = {"chat_id": target_id}
                    if caption:
                        data["caption"] = caption
                    if thread_id:
                        data["message_thread_id"] = int(thread_id)

                    files = {"document": (filename, f, mime_type)}
                    if mime_type.startswith("image/"):
                        method = "sendPhoto"
                        files = {"photo": (filename, f, mime_type)}
                    elif mime_type.startswith("video/"):
                        method = "sendVideo"
                        files = {"video": (filename, f, mime_type)}
                    elif mime_type.startswith("audio/"):
                        method = "sendAudio"
                        files = {"audio": (filename, f, mime_type)}

                    resp = await client.post(
                        f"{TELEGRAM_API_BASE}/bot{token}/{method}",
                        data=data,
                        files=files,
                    )
                    result = resp.json() if resp.content else {}
                    return result.get("result") if result.get("ok") else None
            except Exception:
                logger.exception("Failed to send Telegram file: %s", local_path)
                return None

    @staticmethod
    async def _get_file_path(token: str, file_id: str) -> dict | None:
        import httpx

        try:
            async with httpx.AsyncClient(timeout=httpx.Timeout(10.0)) as client:
                resp = await client.get(
                    f"{TELEGRAM_API_BASE}/bot{token}/getFile",
                    params={"file_id": file_id},
                )
                data = resp.json() if resp.content else {}
                if data.get("ok"):
                    return data.get("result", {})
        except Exception:
            logger.exception("Telegram getFile failed: %s", file_id)
        return None


def _guess_mime(filename: str) -> str:
    ext = filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
    mapping = {
        "jpg": "image/jpeg", "jpeg": "image/jpeg", "png": "image/png",
        "gif": "image/gif", "webp": "image/webp",
        "mp4": "video/mp4", "mov": "video/quicktime",
        "mp3": "audio/mpeg", "ogg": "audio/ogg", "wav": "audio/wav",
        "pdf": "application/pdf",
    }
    return mapping.get(ext, "application/octet-stream")


def _classify_by_extension(filename: str) -> str:
    import os
    from yuxi.channel.utils.file_service import IMAGE_EXTENSIONS, VIDEO_EXTENSIONS

    ext = os.path.splitext(filename)[1].lower()
    if ext in IMAGE_EXTENSIONS:
        return "image"
    if ext in VIDEO_EXTENSIONS:
        return "video"
    return "file"