utils/drive_saver.py

# Save final post-process to Google Drive
import os, json, logging
from typing import Optional

# Check if we're in local mode
IS_LOCAL = os.getenv("IS_LOCAL", "false").lower() == "true"

# Conditional imports for Google Drive (only when not in local mode)
if not IS_LOCAL:
    try:
        from google.oauth2 import service_account
        from googleapiclient.discovery import build
        from googleapiclient.http import MediaFileUpload
        GOOGLE_DRIVE_AVAILABLE = True
    except ImportError:
        GOOGLE_DRIVE_AVAILABLE = False
        # Create dummy classes for when Google Drive is not available
        class service_account:
            class Credentials:
                @staticmethod
                def from_service_account_info(*args, **kwargs):
                    raise ImportError("Google Drive dependencies not available")
        
        class build:
            @staticmethod
            def build(*args, **kwargs):
                raise ImportError("Google Drive dependencies not available")
        
        class MediaFileUpload:
            def __init__(self, *args, **kwargs):
                raise ImportError("Google Drive dependencies not available")
else:
    # Local mode: Google Drive is not available
    GOOGLE_DRIVE_AVAILABLE = False
    # Create dummy classes for local mode
    class service_account:
        class Credentials:
            @staticmethod
            def from_service_account_info(*args, **kwargs):
                raise ImportError("Google Drive not available in local mode")
    
    class build:
        @staticmethod
        def build(*args, **kwargs):
            raise ImportError("Google Drive not available in local mode")
    
    class MediaFileUpload:
        def __init__(self, *args, **kwargs):
            raise ImportError("Google Drive not available in local mode")

from utils.token import get_credentials

logger = logging.getLogger("dsaver")
if not logger.handlers:
    logger.setLevel(logging.INFO)
    fmt = logging.Formatter("[%(levelname)s] %(asctime)s - %(message)s")
    handler = logging.StreamHandler()
    handler.setFormatter(fmt)
    logger.addHandler(handler)

class DriveSaver:
    """Google Drive uploader. Prefers OAuth; optional SA fallback (Shared Drive only)."""

    def __init__(self, default_folder_id: Optional[str] = None):
        self.service = None
        self.folder_id = default_folder_id or os.getenv("GDRIVE_FOLDER_ID")
        self.supports_all_drives = os.getenv("GDRIVE_FOLDER_IS_SHARED", "false").lower() in ("1","true","yes")
        self.allow_sa_fallback = os.getenv("GDRIVE_ALLOW_SA_FALLBACK", "false").lower() in ("1","true","yes")
        
        # Check if we're in local mode
        if IS_LOCAL:
            logger.info("🏠 Local mode: Google Drive integration disabled")
            return
            
        # Check if Google Drive is available
        if not GOOGLE_DRIVE_AVAILABLE:
            logger.warning("⚠️ Google Drive dependencies not available - DriveSaver will be disabled")
            return
            
        if not self.folder_id:
            logger.warning("📁 No GDRIVE_FOLDER_ID set; uploads must provide folder_id explicitly")
        self._initialize_service()

    def _initialize_service(self):
        if IS_LOCAL:
            logger.info("🏠 Local mode: Skipping Google Drive service initialization")
            return
            
        if not GOOGLE_DRIVE_AVAILABLE:
            logger.warning("⚠️ Google Drive dependencies not available - skipping service initialization")
            return
            
        creds = get_credentials()
        if creds:
            logger.info("✅ Using OAuth credentials")
        else:
            # Optional SA fallback — ONLY valid for Shared Drives where SA is a member
            if self.allow_sa_fallback:
                creds_env = os.getenv("GDRIVE_CREDENTIALS_JSON")
                if creds_env:
                    try:
                        info = json.loads(creds_env)
                        if info.get("type") == "service_account":
                            creds = service_account.Credentials.from_service_account_info(
                                info, scopes=["https://www.googleapis.com/auth/drive"]
                            )
                            logger.info("✅ Using Service Account credentials (fallback)")
                            if not self.supports_all_drives:
                                logger.warning("⚠️ SA fallback without Shared Drive mode will likely fail (no quota). "
                                               "Set GDRIVE_FOLDER_IS_SHARED=true and use a Shared Drive folder ID.")
                        else:
                            logger.error("❌ GDRIVE_CREDENTIALS_JSON is not a service account JSON")
                    except Exception as e:
                        logger.error(f"❌ Failed to init Service Account: {e}")
            if not creds:
                logger.error("❌ No valid Google credentials available (OAuth or SA).")
                self.service = None
                return
        # Build Drive service
        self.service = build("drive", "v3", credentials=creds)
        logger.info("✅ Google Drive service initialized")

    def upload_file_to_drive(self, file_path: str, folder_id: Optional[str] = None, mimetype: Optional[str] = None) -> bool:
        if IS_LOCAL:
            logger.info("🏠 Local mode: File upload to Google Drive skipped")
            return False
            
        if not GOOGLE_DRIVE_AVAILABLE:
            logger.warning("⚠️ Google Drive dependencies not available - upload skipped")
            return False
            
        if not self.service:
            logger.error("❌ Drive service not initialized")
            return False
        try:
            target_folder = folder_id or self.folder_id
            name = os.path.basename(file_path)
            media = MediaFileUpload(file_path, mimetype=mimetype or "application/octet-stream")
            metadata = {"name": name, "parents": [target_folder]}
            req = self.service.files().create(
                body=metadata,
                media_body=media,
                fields="id",
                supportsAllDrives=self.supports_all_drives
            )
            req.execute()
            logger.info(f"✅ Uploaded '{name}' to Drive (folder: {target_folder})")
            return True
        except Exception as e:
            logger.error(f"❌ Drive upload failed: {e}")
            return False

    def is_service_available(self) -> bool:
        return not IS_LOCAL and GOOGLE_DRIVE_AVAILABLE and self.service is not None

    def set_folder_id(self, folder_id: str):
        self.folder_id = folder_id
        logger.info(f"📁 Default folder ID updated: {folder_id}")