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

from __future__ import annotations

import logging
import os

import httpx

from .constants import ENDPOINT_GET_UPLOAD_URL, ENDPOINT_SEND_MESSAGE
from .types import FlockOutboundResult, FlockUploadResult
from .utils import create_http_client, flock_api_call

logger = logging.getLogger(__name__)


async def upload_file(
    file_path: str,
    file_name: str,
    content_type: str,
    target_id: str,
    bot_token: str,
    client: httpx.AsyncClient | None = None,
) -> FlockUploadResult:
    own_client = client is None
    if own_client:
        client = create_http_client()

    try:
        file_size = os.path.getsize(file_path)

        upload_url_resp = await _get_upload_url(client, bot_token, target_id, file_name, file_size, content_type)
        if not upload_url_resp.success:
            return upload_url_resp

        upload_url = upload_url_resp.upload_url
        file_id = upload_url_resp.file_id

        if upload_url:
            await _upload_file_content(client, upload_url, file_path, file_name, content_type)

        return FlockUploadResult(success=True, file_id=file_id, upload_url=upload_url)

    except Exception as e:
        logger.error("Flock file upload failed: %s", e)
        return FlockUploadResult(success=False, error=str(e))
    finally:
        if own_client:
            await client.aclose()


async def _get_upload_url(
    client: httpx.AsyncClient,
    bot_token: str,
    target_id: str,
    file_name: str,
    file_size: int,
    content_type: str,
) -> FlockUploadResult:
    payload = {
        "to": target_id,
        "fileName": file_name,
        "fileSize": file_size,
        "contentType": content_type,
    }
    try:
        result = await flock_api_call(client, ENDPOINT_GET_UPLOAD_URL, bot_token, payload)
        upload_url = result.get("uploadUrl", "")
        file_id = result.get("fileId", "")
        return FlockUploadResult(success=True, file_id=file_id, upload_url=upload_url)
    except Exception as e:
        logger.error("Flock getUploadUrl failed: %s", e)
        return FlockUploadResult(success=False, error=str(e))


async def _upload_file_content(
    client: httpx.AsyncClient,
    upload_url: str,
    file_path: str,
    file_name: str,
    content_type: str,
) -> None:
    with open(file_path, "rb") as f:
        files = {"file": (file_name, f, content_type)}
        upload_client = httpx.AsyncClient(timeout=60.0)
        try:
            resp = await upload_client.post(upload_url, files=files)
            if resp.status_code not in (200, 201, 204):
                logger.warning("Flock file upload response: HTTP %d", resp.status_code)
                raise RuntimeError(f"File upload failed: HTTP {resp.status_code}")
        finally:
            await upload_client.aclose()


async def send_image(
    target_id: str,
    image_url: str,
    title: str = "",
    *,
    thread_id: str | None = None,
    bot_token: str,
    client: httpx.AsyncClient | None = None,
) -> FlockOutboundResult:
    payload: dict = {
        "to": target_id,
        "text": title or "Image",
        "attachments": [
            {
                "title": title or "Image",
                "views": {
                    "image": {
                        "original": {"src": image_url},
                        "thumbnail": {"src": image_url},
                    }
                },
            }
        ],
    }
    if thread_id:
        payload["threadId"] = thread_id

    own_client = client is None
    if own_client:
        client = create_http_client()

    try:
        result = await flock_api_call(client, ENDPOINT_SEND_MESSAGE, bot_token, payload)
        return FlockOutboundResult(success=True, message_uid=result.get("uid", ""))
    except Exception as e:
        logger.error("Flock image send failed: %s", e)
        retryable = getattr(e, "retryable", False)
        return FlockOutboundResult(success=False, error=str(e), retryable=retryable)
    finally:
        if own_client:
            await client.aclose()


async def send_file(
    target_id: str,
    file_path: str,
    file_name: str,
    *,
    thread_id: str | None = None,
    bot_token: str,
    client: httpx.AsyncClient | None = None,
) -> FlockOutboundResult:
    import mimetypes

    content_type, _ = mimetypes.guess_type(file_name)
    if not content_type:
        content_type = "application/octet-stream"

    own_client = client is None
    if own_client:
        client = create_http_client()

    try:
        upload_result = await upload_file(file_path, file_name, content_type, target_id, bot_token, client)
        if not upload_result.success:
            return FlockOutboundResult(
                success=False,
                error=f"File upload failed: {upload_result.error}",
            )

        payload: dict = {
            "to": target_id,
            "text": file_name,
            "attachments": [
                {
                    "title": file_name,
                    "fileId": upload_result.file_id,
                }
            ],
        }
        if thread_id:
            payload["threadId"] = thread_id

        result = await flock_api_call(client, ENDPOINT_SEND_MESSAGE, bot_token, payload)
        return FlockOutboundResult(success=True, message_uid=result.get("uid", ""))
    except Exception as e:
        logger.error("Flock file send failed: %s", e)
        retryable = getattr(e, "retryable", False)
        return FlockOutboundResult(success=False, error=str(e), retryable=retryable)
    finally:
        if own_client:
            await client.aclose()