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

from __future__ import annotations

import logging
import os
import tempfile

import httpx

logger = logging.getLogger(__name__)

ZOOM_FILE_UPLOAD_URL = "/chat/users/{user_id}/files"
ZOOM_FILE_DOWNLOAD_URL = "/chat/users/{user_id}/messages/{message_id}/files/{file_id}"


async def upload_file(
    client: httpx.AsyncClient,
    token: str,
    bot_user_id: str,
    file_path: str,
    media_type: str = "application/octet-stream",
    target_channel_id: str | None = None,
    target_contact: str | None = None,
) -> str | None:
    if file_path.startswith("file://"):
        file_path = file_path[7:]

    if file_path.startswith(("http://", "https://")):
        file_path = await _download_to_tmp(client, file_path)

    if not os.path.exists(file_path):
        logger.error("Media file not found: %s", file_path)
        return None

    filename = os.path.basename(file_path)
    url = ZOOM_FILE_UPLOAD_URL.format(user_id=bot_user_id)

    with open(file_path, "rb") as f:
        resp = await client.post(
            url,
            files={"file": (filename, f, media_type)},
            headers={"Authorization": f"Bearer {token}"},
        )

    if resp.status_code not in (200, 201):
        logger.error("Zoom file upload failed: status=%s, body=%s", resp.status_code, resp.text)
        return None

    file_id = resp.json().get("id", "")
    if not file_id:
        logger.error("Zoom file upload returned no file_id")
        return None

    logger.info("Zoom file uploaded: file_id=%s, name=%s", file_id, filename)
    return file_id


async def download_file(
    client: httpx.AsyncClient,
    token: str,
    bot_user_id: str,
    message_id: str,
    file_id: str,
    save_dir: str | None = None,
) -> str | None:
    url = ZOOM_FILE_DOWNLOAD_URL.format(user_id=bot_user_id, message_id=message_id, file_id=file_id)

    resp = await client.get(
        url,
        headers={"Authorization": f"Bearer {token}"},
        follow_redirects=True,
    )

    if resp.status_code != 200:
        logger.error("Zoom file download failed: status=%s", resp.status_code)
        return None

    content_disposition = resp.headers.get("content-disposition", "")
    filename = _extract_filename(content_disposition) or f"{file_id}.bin"

    if save_dir:
        os.makedirs(save_dir, exist_ok=True)
        save_path = os.path.join(save_dir, filename)
    else:
        ext = _guess_extension(resp.content)
        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=ext)
        save_path = tmp.name
        tmp.close()

    with open(save_path, "wb") as f:
        f.write(resp.content)

    logger.info("Zoom file downloaded: file_id=%s -> %s", file_id, save_path)
    return save_path


async def send_file_message(
    client: httpx.AsyncClient,
    token: str,
    bot_user_id: str,
    file_id: str,
    target_id: str,
    caption: str = "",
    reply_to_id: str | None = None,
) -> str | None:
    url = f"/chat/users/{bot_user_id}/messages"

    body: dict = {
        "message": caption or "[附件]",
        "file_ids": [file_id],
    }
    if target_id.startswith("ch_"):
        body["to_channel"] = target_id
    else:
        body["to_contact"] = target_id
    if reply_to_id:
        body["reply_main_message_id"] = reply_to_id

    resp = await client.post(
        url,
        json=body,
        headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
    )

    if resp.status_code in (200, 201):
        return resp.json().get("id", "")

    logger.error("Zoom file message send failed: status=%s, body=%s", resp.status_code, resp.text)
    return None


async def _download_to_tmp(client: httpx.AsyncClient, url: str) -> str:
    resp = await client.get(url, follow_redirects=True, timeout=60.0)
    resp.raise_for_status()
    data = resp.content
    ext = _guess_extension(data)

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=ext)
    tmp.write(data)
    tmp.close()
    return tmp.name


def _extract_filename(content_disposition: str) -> str:
    if not content_disposition:
        return ""
    for part in content_disposition.split(";"):
        part = part.strip()
        if part.startswith("filename="):
            return part[len("filename=") :].strip('"')
    return ""


def _guess_extension(data: bytes) -> str:
    if data[:4] == b"\x89PNG":
        return ".png"
    if data[:2] == b"\xff\xd8":
        return ".jpg"
    if data[:4] == b"GIF8":
        return ".gif"
    if data[:4] == b"RIFF" and len(data) > 11 and data[8:12] == b"WEBP":
        return ".webp"
    if data[:5] == b"%PDF-":
        return ".pdf"
    return ".bin"