ForcePilot/backend/package/yuxi/channels/adapters/googlechat/media.py

from __future__ import annotations

import io
import json
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from googleapiclient.discovery import Resource

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

_MAX_DOWNLOAD_SIZE_BYTES = 100 * 1024 * 1024


async def upload_image(
    chat_service: Resource,
    chat_id: str,
    image_data: bytes,
    filename: str = "image.png",
    caption: str = "",
) -> DeliveryResult:
    from googleapiclient.http import MediaIoBaseUpload

    try:
        media = MediaIoBaseUpload(
            io.BytesIO(image_data),
            mimetype="image/png",
            resumable=True,
        )
        body = {"text": caption or " "}
        result = chat_service.spaces().messages().create(parent=chat_id, body=body, media_body=media).execute()
        return DeliveryResult(
            success=True,
            message_id=result.get("name", ""),
        )
    except Exception as e:
        logger.error(f"upload_image failed: {e}")
        return DeliveryResult(success=False, error=str(e))


async def upload_file(
    chat_service: Resource,
    chat_id: str,
    file_data: bytes,
    filename: str,
    mime_type: str = "application/octet-stream",
    caption: str = "",
) -> DeliveryResult:
    from googleapiclient.http import MediaIoBaseUpload

    try:
        media = MediaIoBaseUpload(
            io.BytesIO(file_data),
            mimetype=mime_type,
            resumable=True,
        )
        body = {"text": caption or " "}
        result = chat_service.spaces().messages().create(parent=chat_id, body=body, media_body=media).execute()
        return DeliveryResult(
            success=True,
            message_id=result.get("name", ""),
        )
    except Exception as e:
        logger.error(f"upload_file failed: {e}")
        return DeliveryResult(success=False, error=str(e))


async def download_media(chat_service: Resource, file_id: str) -> bytes:
    try:
        file_info = json.loads(file_id)
    except (json.JSONDecodeError, TypeError):
        raise DeliveryFailedError(f"Invalid file_id format: {file_id}")

    resource_name = file_info.get("name", file_info.get("resource_name", ""))
    if not resource_name:
        raise DeliveryFailedError("file_id missing attachment resource name")

    try:
        attachment = chat_service.spaces().messages().attachments().get(name=resource_name).execute()
    except Exception as e:
        raise DeliveryFailedError(f"Failed to get attachment metadata: {e}")

    drive_data = attachment.get("driveDataRef", {})
    if drive_data and drive_data.get("driveFileId"):
        try:
            return await _download_drive_file(chat_service, drive_data["driveFileId"])
        except Exception as e:
            raise DeliveryFailedError(f"Drive file download failed (file_id={drive_data.get('driveFileId')}): {e}")

    try:
        creds = _get_credentials(chat_service)
        token = await _get_access_token(creds)

        import httpx

        url = f"https://chat.googleapis.com/v1/{resource_name}?alt=media"
        async with httpx.AsyncClient(timeout=httpx.Timeout(60)) as client:
            resp = await client.get(
                url,
                headers={"Authorization": f"Bearer {token}"},
                follow_redirects=True,
            )
            if resp.status_code == 200:
                content_length = int(resp.headers.get("content-length", 0))
                if content_length > _MAX_DOWNLOAD_SIZE_BYTES:
                    raise DeliveryFailedError(
                        f"File too large: {content_length} bytes (max {_MAX_DOWNLOAD_SIZE_BYTES})"
                    )
                return resp.content
            raise DeliveryFailedError(f"Media download returned HTTP {resp.status_code}")
    except DeliveryFailedError:
        raise
    except Exception as e:
        logger.warning(f"Direct media download failed, falling back: {e}")

    raise DeliveryFailedError(f"Unable to download media for attachment: {resource_name}")


async def _download_drive_file(chat_service: Resource, drive_file_id: str) -> bytes:
    from googleapiclient.discovery import build
    from googleapiclient.http import MediaIoBaseDownload

    creds = _get_credentials(chat_service)
    drive_service = build("drive", "v3", credentials=creds, cache_discovery=False)

    request = drive_service.files().get_media(fileId=drive_file_id)
    buf = io.BytesIO()
    downloader = MediaIoBaseDownload(buf, request)
    done = False
    while not done:
        _, done = downloader.next_chunk()
    return buf.getvalue()


def _get_credentials(chat_service: Resource):
    http = getattr(chat_service, "_http", None)
    if http is None:
        raise DeliveryFailedError("No HTTP transport on chat_service")
    creds = getattr(http, "credentials", None)
    if creds is None:
        raise DeliveryFailedError("No credentials on chat_service HTTP transport")
    return creds


async def _get_access_token(credentials) -> str:
    from google.auth.transport.requests import Request as GARequest

    if credentials.valid:
        return credentials.token
    credentials.refresh(GARequest())
    return credentials.token