google.py

import os
import pickle
from google.auth.transport.requests import Request
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
from googleapiclient.http import MediaFileUpload, MediaIoBaseUpload
import io
from pathlib import Path
import logging

logger = logging.getLogger(__name__)

# --- CONFIGURATION ---
# Get credentials from environment variables
CLIENT_ID = os.getenv("GOOGLE_CLIENT_ID")
CLIENT_SECRET = os.getenv("GOOGLE_CLIENT_SECRET")

# Google Drive folder IDs
EXPORTS_FOLDER_ID = "1k5iP4egzLrGJwnHkMhxt9bAkaCiieojO"  # For Excel exports
IMAGES_FOLDER_ID = "1gd280IqcAzpAFTPeYsZjoBUOU9S7Zx3c"   # For business card images

# Scopes define the level of access you are requesting.
SCOPES = ['https://www.googleapis.com/auth/drive.file']
TOKEN_PICKLE_FILE = 'token.pickle'

def get_drive_service():
    """Authenticates with Google and returns a Drive service object."""
    creds = None
    # The file token.pickle stores the user's access and refresh tokens.
    if os.path.exists(TOKEN_PICKLE_FILE):
        with open(TOKEN_PICKLE_FILE, 'rb') as token:
            creds = pickle.load(token)
            
    # If there are no (valid) credentials available, let the user log in.
    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            creds.refresh(Request())
        else:
            if not CLIENT_ID or not CLIENT_SECRET:
                raise ValueError("GOOGLE_CLIENT_ID and GOOGLE_CLIENT_SECRET environment variables are required")
                
            # Use client_config dictionary instead of a client_secret.json file
            client_config = {
                "installed": {
                    "client_id": CLIENT_ID,
                    "client_secret": CLIENT_SECRET,
                    "auth_uri": "https://accounts.google.com/o/oauth2/auth",
                    "token_uri": "https://oauth2.googleapis.com/token",
                    "redirect_uris": ["http://localhost"]
                }
            }
            flow = InstalledAppFlow.from_client_config(client_config, SCOPES)
            creds = flow.run_local_server(port=0)
            
        # Save the credentials for the next run
        with open(TOKEN_PICKLE_FILE, 'wb') as token:
            pickle.dump(creds, token)
            
    return build('drive', 'v3', credentials=creds)

def upload_file_to_drive(service, file_path=None, file_data=None, filename=None, folder_id=None, mimetype='application/octet-stream'):
    """
    Uploads a file to a specific folder in Google Drive.
    
    Args:
        service: Google Drive service object
        file_path: Path to local file (for file uploads)
        file_data: Bytes data (for in-memory uploads)
        filename: Name for the file in Drive
        folder_id: ID of the target folder
        mimetype: MIME type of the file
    
    Returns:
        dict: File information (id, webViewLink) or None if failed
    """
    try:
        if file_path and os.path.exists(file_path):
            # Upload from local file
            if not filename:
                filename = os.path.basename(file_path)
            media = MediaFileUpload(file_path, mimetype=mimetype, resumable=True)
            logger.info(f"Uploading file from path: {file_path}")
        elif file_data and filename:
            # Upload from bytes data
            file_io = io.BytesIO(file_data)
            media = MediaIoBaseUpload(file_io, mimetype=mimetype, resumable=True)
            logger.info(f"Uploading file from memory: {filename}")
        else:
            logger.error("Either file_path or (file_data + filename) must be provided")
            return None

        # Define the file's metadata
        file_metadata = {
            'name': filename,
            'parents': [folder_id] if folder_id else []
        }

        logger.info(f"Uploading '{filename}' to Google Drive folder {folder_id}")

        # Execute the upload request
        file = service.files().create(
            body=file_metadata,
            media_body=media,
            fields='id, webViewLink, name'
        ).execute()

        logger.info(f"✅ File uploaded successfully!")
        logger.info(f"   File ID: {file.get('id')}")
        logger.info(f"   File Name: {file.get('name')}")
        logger.info(f"   View Link: {file.get('webViewLink')}")

        return {
            'id': file.get('id'),
            'name': file.get('name'),
            'webViewLink': file.get('webViewLink')
        }

    except Exception as e:
        logger.error(f"Failed to upload file to Google Drive: {e}")
        return None

def upload_excel_to_exports_folder(service, file_path=None, file_data=None, filename=None):
    """Upload Excel file to the exports folder."""
    return upload_file_to_drive(
        service, 
        file_path=file_path, 
        file_data=file_data, 
        filename=filename, 
        folder_id=EXPORTS_FOLDER_ID,
        mimetype='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
    )

def upload_image_to_images_folder(service, file_path=None, file_data=None, filename=None, mimetype='image/png'):
    """Upload image file to the images folder."""
    return upload_file_to_drive(
        service, 
        file_path=file_path, 
        file_data=file_data, 
        filename=filename, 
        folder_id=IMAGES_FOLDER_ID,
        mimetype=mimetype
    )

def list_files_in_folder(service, folder_id, max_results=100):
    """List files in a specific Google Drive folder."""
    try:
        query = f"'{folder_id}' in parents"
        results = service.files().list(
            q=query,
            maxResults=max_results,
            fields="files(id, name, size, createdTime, webViewLink)"
        ).execute()
        
        files = results.get('files', [])
        logger.info(f"Found {len(files)} files in folder {folder_id}")
        return files
    except Exception as e:
        logger.error(f"Failed to list files in folder {folder_id}: {e}")
        return []

if __name__ == '__main__':
    # Test the Google Drive connection
    try:
        drive_service = get_drive_service()
        logger.info("Google Drive service initialized successfully")
        
        # List files in both folders to verify access
        exports_files = list_files_in_folder(drive_service, EXPORTS_FOLDER_ID)
        images_files = list_files_in_folder(drive_service, IMAGES_FOLDER_ID)
        
        print(f"Exports folder contains {len(exports_files)} files")
        print(f"Images folder contains {len(images_files)} files")
        
    except Exception as e:
        logger.error(f"Failed to initialize Google Drive: {e}")