ForcePilot/backend/package/yuxi/channels/adapters/qqbot/media_upload.py

from __future__ import annotations

import asyncio

import aiohttp
from cachetools import LRUCache

from yuxi.channels.exceptions import DeliveryFailedError
from yuxi.utils.logging_config import logger

_MAX_FILE_SIZE_MB = 100
FILE_TYPE_IMAGE = "1"
FILE_TYPE_VOICE = "2"
FILE_TYPE_VIDEO = "3"
FILE_TYPE_FILE = "4"


def validate_media_size(data: bytes, max_size_mb: int = _MAX_FILE_SIZE_MB, label: str = "media") -> None:
    max_bytes = max_size_mb * 1024 * 1024
    actual_size = len(data)
    if actual_size > max_bytes:
        raise DeliveryFailedError(f"{label} size {actual_size / 1024 / 1024:.1f}MB exceeds limit of {max_size_mb}MB")


async def upload_media(
    media_data: bytes,
    token: str,
    http_client: aiohttp.ClientSession | None = None,
    filename: str = "media",
    file_type: str = FILE_TYPE_FILE,
    group_openid: str | None = None,
    sandbox: bool = False,
) -> str:
    api_base = "https://sandbox.api.sgroup.qq.com" if sandbox else "https://api.sgroup.qq.com"

    headers = {"Authorization": f"QQBot {token}"}

    content_type_map = {
        FILE_TYPE_IMAGE: "image/png",
        FILE_TYPE_VOICE: "audio/mpeg",
        FILE_TYPE_VIDEO: "video/mp4",
        FILE_TYPE_FILE: "application/octet-stream",
    }
    mime_type = content_type_map.get(file_type, content_type_map[FILE_TYPE_FILE])

    form = aiohttp.FormData()
    form.add_field("file", media_data, filename=filename, content_type=mime_type)
    form.add_field("file_type", file_type)

    url = f"{api_base}/v2/groups/{group_openid}/files" if group_openid else f"{api_base}/v2/users/@me/files"

    async def _do_upload(client: aiohttp.ClientSession) -> str:
        last_error = None
        for attempt in range(3):
            try:
                async with client.post(url, headers=headers, data=form) as resp:
                    if resp.status != 200:
                        raise DeliveryFailedError(f"Media upload failed: HTTP {resp.status}")
                    result = await resp.json()
                    file_id = result.get("file_uuid", "") or result.get("file_info", "")
                    logger.debug(f"[QQBot] Media uploaded, file_type={file_type}, file_id={file_id}")
                    return file_id
            except DeliveryFailedError:
                raise
            except Exception as e:
                last_error = e
                if attempt < 2:
                    delay = 2**attempt
                    logger.warning("[QQBot] Media upload retry %d/%d after %.1fs: %s", attempt + 1, 3, delay, e)
                    await asyncio.sleep(delay)
        raise DeliveryFailedError(f"Media upload failed after 3 attempts: {last_error}")

    if http_client:
        return await _do_upload(http_client)
    else:
        async with aiohttp.ClientSession() as session:
            return await _do_upload(session)


async def upload_image(
    image_data: bytes,
    token: str,
    http_client: aiohttp.ClientSession | None = None,
    filename: str = "image.png",
    group_openid: str | None = None,
    sandbox: bool = False,
) -> str:
    return await upload_media(
        image_data,
        token,
        http_client=http_client,
        filename=filename,
        file_type=FILE_TYPE_IMAGE,
        group_openid=group_openid,
        sandbox=sandbox,
    )


async def download_media(
    file_id: str,
    token: str,
    http_client: aiohttp.ClientSession | None = None,
    sandbox: bool = False,
) -> bytes:
    api_base = "https://sandbox.api.sgroup.qq.com" if sandbox else "https://api.sgroup.qq.com"

    headers = {"Authorization": f"QQBot {token}"}

    async def _do_download(client: aiohttp.ClientSession) -> bytes:
        async with client.get(
            f"{api_base}/v2/users/@me/files/{file_id}",
            headers=headers,
        ) as resp:
            if resp.status != 200:
                raise DeliveryFailedError(f"Media download failed: HTTP {resp.status}")
            return await resp.read()

    if http_client:
        return await _do_download(http_client)
    else:
        async with aiohttp.ClientSession() as session:
            return await _do_download(session)


def build_media_payload(
    chat_id: str,
    file_id: str,
    content: str = "",
    msg_type: int = 7,
) -> dict:
    from .constants import DM_CHAT_PREFIX, GROUP_CHAT_PREFIX

    payload: dict = {
        "msg_type": msg_type,
    }
    if msg_type == 1:
        payload["file_image"] = file_id
    elif msg_type == 4:
        payload["file"] = file_id
    else:
        payload["media"] = {"file_info": file_id}

    if content:
        payload["content"] = content[:2000]

    if chat_id.startswith(GROUP_CHAT_PREFIX):
        payload["group_openid"] = chat_id.replace(GROUP_CHAT_PREFIX, "")
    elif not chat_id.startswith(DM_CHAT_PREFIX):
        payload["channel_id"] = chat_id

    return payload


async def download_image(
    file_id: str,
    token: str,
    http_client: aiohttp.ClientSession | None = None,
    sandbox: bool = False,
) -> bytes:
    return await download_media(file_id, token, http_client, sandbox)


_CHUNK_SIZE = 5 * 1024 * 1024
_UPLOAD_CACHE = LRUCache(maxsize=1000)


def _make_cache_key(data: bytes) -> str:
    import hashlib

    return hashlib.sha256(data).hexdigest()


async def upload_media_cached(
    media_data: bytes,
    token: str,
    http_client: aiohttp.ClientSession | None = None,
    filename: str = "media",
    file_type: str = FILE_TYPE_FILE,
    group_openid: str | None = None,
    sandbox: bool = False,
    use_cache: bool = True,
) -> str:
    if use_cache:
        cache_key = _make_cache_key(media_data)
        cached = _UPLOAD_CACHE.get(cache_key)
        if cached:
            logger.debug("Media upload: cache hit for %s", filename)
            return cached

    file_id = await upload_media(
        media_data,
        token,
        http_client=http_client,
        filename=filename,
        file_type=file_type,
        group_openid=group_openid,
        sandbox=sandbox,
    )

    if use_cache and file_id:
        cache_key = _make_cache_key(media_data)
        _UPLOAD_CACHE[cache_key] = file_id

    return file_id


def clear_upload_cache() -> None:
    _UPLOAD_CACHE.clear()
    logger.debug("Media upload cache cleared")


async def upload_media_chunked(
    media_data: bytes,
    token: str,
    http_client: aiohttp.ClientSession | None = None,
    filename: str = "media",
    file_type: str = FILE_TYPE_FILE,
    group_openid: str | None = None,
    sandbox: bool = False,
    chunk_size: int = _CHUNK_SIZE,
) -> str:
    if len(media_data) <= chunk_size:
        return await upload_media(
            media_data,
            token,
            http_client=http_client,
            filename=filename,
            file_type=file_type,
            group_openid=group_openid,
            sandbox=sandbox,
        )

    import math

    total_chunks = math.ceil(len(media_data) / chunk_size)
    api_base = "https://sandbox.api.sgroup.qq.com" if sandbox else "https://api.sgroup.qq.com"
    headers = {"Authorization": f"QQBot {token}"}

    content_type_map = {
        FILE_TYPE_IMAGE: "image/png",
        FILE_TYPE_VOICE: "audio/mpeg",
        FILE_TYPE_VIDEO: "video/mp4",
        FILE_TYPE_FILE: "application/octet-stream",
    }
    mime_type = content_type_map.get(file_type, content_type_map[FILE_TYPE_FILE])

    async def _do_chunked(client: aiohttp.ClientSession) -> str:
        init_url = f"{api_base}/v2/users/@me/files/chunked"
        if group_openid:
            init_url = f"{api_base}/v2/groups/{group_openid}/files/chunked"

        init_payload = {
            "filename": filename,
            "file_type": int(file_type),
            "total_size": len(media_data),
            "chunk_size": chunk_size,
            "total_chunks": total_chunks,
        }

        async with client.post(
            init_url,
            headers=headers,
            json=init_payload,
        ) as resp:
            if resp.status not in (200, 201):
                raise DeliveryFailedError(f"Chunked upload init failed: HTTP {resp.status}")
            init_data = await resp.json()
            upload_id = init_data.get("upload_id", "")

        if not upload_id:
            raise DeliveryFailedError("Chunked upload: no upload_id returned")

        for i in range(total_chunks):
            start = i * chunk_size
            end = min(start + chunk_size, len(media_data))
            chunk = media_data[start:end]

            chunk_url = f"{api_base}/v2/users/@me/files/chunked/{upload_id}"
            if group_openid:
                chunk_url = f"{api_base}/v2/groups/{group_openid}/files/chunked/{upload_id}"

            form = aiohttp.FormData()
            form.add_field("chunk", chunk, filename=f"{filename}.chunk{i}", content_type=mime_type)
            form.add_field("chunk_index", str(i))

            async with client.post(chunk_url, headers=headers, data=form) as resp:
                if resp.status not in (200, 201):
                    raise DeliveryFailedError(f"Chunked upload part {i + 1}/{total_chunks} failed: HTTP {resp.status}")

        complete_url = f"{api_base}/v2/users/@me/files/chunked/{upload_id}/complete"
        if group_openid:
            complete_url = f"{api_base}/v2/groups/{group_openid}/files/chunked/{upload_id}/complete"

        async with client.post(complete_url, headers=headers) as resp:
            if resp.status != 200:
                raise DeliveryFailedError(f"Chunked upload complete failed: HTTP {resp.status}")
            result = await resp.json()
            return result.get("file_uuid", "") or result.get("file_info", "")

    if http_client:
        return await _do_chunked(http_client)
    else:
        async with aiohttp.ClientSession() as session:
            return await _do_chunked(session)


async def upload_media_from_url(
    url: str,
    token: str,
    http_client: aiohttp.ClientSession | None = None,
    filename: str = "media",
    file_type: str = FILE_TYPE_FILE,
    group_openid: str | None = None,
    sandbox: bool = False,
) -> str:
    async def _do_url_upload(client: aiohttp.ClientSession) -> str:
        api_base = "https://sandbox.api.sgroup.qq.com" if sandbox else "https://api.sgroup.qq.com"
        headers = {"Authorization": f"QQBot {token}"}

        endpoint = f"{api_base}/v2/users/@me/files/url"
        if group_openid:
            endpoint = f"{api_base}/v2/groups/{group_openid}/files/url"

        payload = {
            "url": url,
            "file_type": int(file_type),
        }

        async with client.post(endpoint, headers=headers, json=payload) as resp:
            if resp.status not in (200, 201):
                raise DeliveryFailedError(f"URL upload failed: HTTP {resp.status}")
            result = await resp.json()
            return result.get("file_uuid", "") or result.get("file_info", "")

    if http_client:
        return await _do_url_upload(http_client)
    else:
        async with aiohttp.ClientSession() as session:
            return await _do_url_upload(session)