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.utils.logging_config import logger _MAX_DOWNLOAD_SIZE_BYTES = 100 * 1024 * 1024 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