src/drive_uploader.py

"""
Google Drive Auto-Uploader using OAuth2 refresh token.
Works with personal Google accounts (not just Workspace).
"""
import os
import json
from google.oauth2.credentials import Credentials
from google.auth.transport.requests import Request
from googleapiclient.discovery import build
from googleapiclient.http import MediaFileUpload, MediaIoBaseDownload
import io

# Scopes required for Drive upload
SCOPES = ['https://www.googleapis.com/auth/drive']

def get_drive_service():
    """
    Creates a Google Drive API service using OAuth2 refresh token.
    Requires 3 env vars: GOOGLE_OAUTH_CLIENT_ID, GOOGLE_OAUTH_CLIENT_SECRET, GOOGLE_DRIVE_REFRESH_TOKEN
    """
    client_id = os.getenv("GOOGLE_OAUTH_CLIENT_ID")
    client_secret = os.getenv("GOOGLE_OAUTH_CLIENT_SECRET")
    refresh_token = os.getenv("GOOGLE_DRIVE_REFRESH_TOKEN")
    
    if not all([client_id, client_secret, refresh_token]):
        missing = []
        if not client_id: missing.append("GOOGLE_OAUTH_CLIENT_ID")
        if not client_secret: missing.append("GOOGLE_OAUTH_CLIENT_SECRET")
        if not refresh_token: missing.append("GOOGLE_DRIVE_REFRESH_TOKEN")
        print(f"ERROR: Missing env vars: {', '.join(missing)}", flush=True)
        return None
    
    try:
        creds = Credentials(
            token=None,
            refresh_token=refresh_token,
            client_id=client_id,
            client_secret=client_secret,
            token_uri="https://oauth2.googleapis.com/token",
            scopes=SCOPES
        )
        # Refresh the access token
        creds.refresh(Request())
        
        service = build('drive', 'v3', credentials=creds)
        print("✅ Google Drive connected (OAuth2)!", flush=True)
        return service
    except Exception as e:
        print(f"❌ Failed to connect to Google Drive: {e}", flush=True)
        return None

def upload_to_drive(file_path, folder_id=None):
    """
    Uploads a file to Google Drive using the user's own account.
    
    Args:
        file_path: Path to the file to upload
        folder_id: Google Drive folder ID
    
    Returns:
        The file URL on Google Drive, or None if upload failed
    """
    if not folder_id:
        folder_id = os.getenv("GOOGLE_DRIVE_FOLDER_ID")
    
    if not folder_id:
        print("ERROR: GOOGLE_DRIVE_FOLDER_ID not set!", flush=True)
        return None
    
    service = get_drive_service()
    if not service:
        return None
    
    filename = os.path.basename(file_path)
    file_size_mb = os.path.getsize(file_path) / (1024 * 1024)
    
    print(f"📤 Uploading {filename} ({file_size_mb:.1f} MB) to Google Drive...", flush=True)
    
    file_metadata = {
        'name': filename,
        'parents': [folder_id]
    }
    
    media = MediaFileUpload(
        file_path,
        mimetype='video/mp4',
        resumable=True,
        chunksize=10 * 1024 * 1024  # 10MB chunks
    )
    
    try:
        file = service.files().create(
            body=file_metadata,
            media_body=media,
            fields='id, webViewLink'
        ).execute()
        
        file_id = file.get('id')
        file_link = file.get('webViewLink', f"https://drive.google.com/file/d/{file_id}/view")
        
        print(f"✅ Upload SUCCESS: {filename}", flush=True)
        print(f"📎 Link: {file_link}", flush=True)
        
        return file_link
        
    except Exception as e:
        print(f"❌ Upload FAILED: {e}", flush=True)
        return None

def list_files(folder_id):
    """
    Lists files in a specific Google Drive folder.
    """
    service = get_drive_service()
    if not service:
        return []

    try:
        results = service.files().list(
            q=f"'{folder_id}' in parents and trashed=false",
            fields="files(id, name, mimeType)",
            pageSize=100
        ).execute()
        files = results.get('files', [])
        return files
    except Exception as e:
        print(f"❌ Failed to list files: {e}", flush=True)
        return []

def download_file(file_id, destination_path):
    """
    Downloads a file from Google Drive to a local path.
    """
    service = get_drive_service()
    if not service:
        return False

    try:
        request = service.files().get_media(fileId=file_id)
        
        print(f"⬇️ Downloading file to {destination_path}...", flush=True)
        
        with open(destination_path, 'wb') as f:
            downloader = MediaIoBaseDownload(f, request)
            done = False
            while done is False:
                status, done = downloader.next_chunk()
                # print(f"Download {int(status.progress() * 100)}%.", flush=True)
        
        print(f"✅ Download complete: {destination_path}", flush=True)
        return True
    except Exception as e:
        print(f"❌ Download FAILED: {e}", flush=True)
        return False