# app.py
import os
import requests
import json
from simple_salesforce import Salesforce
from google.oauth2 import service_account
from googleapiclient.discovery import build
from googleapiclient.http import MediaIoBaseDownload
import tempfile

# Google Drive Authentication Setup
SCOPES = ['https://www.googleapis.com/auth/drive']
SERVICE_ACCOUNT_FILE = 'service_account.json'  # Ensure this file is uploaded to the Hugging Face space
credentials = service_account.Credentials.from_service_account_file(
    SERVICE_ACCOUNT_FILE, scopes=SCOPES)
drive_service = build('drive', 'v3', credentials=credentials)

# Salesforce Authentication
sf = Salesforce(
    username='bharat@sathkrutha.sandbox',
    password='Mahadev@9',
    security_token='q67QF1yG0qBTSaiEfDrRrIEx',
    domain='login'
)

# Using a temporary directory to save retrieved files
temp_dir = tempfile.mkdtemp()

# Function to retrieve all files from the Salesforce custom object
def get_all_salesforce_files(custom_object_api_name):
    # Step 1: Query all record IDs from the custom object, including the Google Drive File ID field
    query_all_records = f"SELECT Id, Google_Fille_ID__c FROM {custom_object_api_name} WHERE Google_Fille_ID__c != NULL"
    all_records = sf.query(query_all_records)
    records = all_records['records']
    
    return records

# Function to download a file from Google Drive and save it to a temporary location
def download_google_drive_file(file_id, file_name):
    try:
        request = drive_service.files().get_media(fileId=file_id)
        file_path = os.path.join(temp_dir, file_name)
        with open(file_path, 'wb') as f:
            downloader = MediaIoBaseDownload(f, request)
            done = False
            while not done:
                status, done = downloader.next_chunk()
                if status:
                    print(f"Download progress: {int(status.progress() * 100)}%")
        return file_path
    except Exception as e:
        print(f"Error downloading file {file_name}: {e}")
        return None

# Function to save all files linked to records from the Salesforce custom object
def save_salesforce_files(custom_object_api_name):
    records = get_all_salesforce_files(custom_object_api_name)
    for record in records:
        google_drive_file_id = record['Google_Drive_File_ID__c']
        record_id = record['Id']
        
        # Check if the Google Drive File ID exists
        if google_drive_file_id:
            # Get the file metadata from Google Drive to obtain the filename
            try:
                file_metadata = drive_service.files().get(fileId=google_drive_file_id, fields='name').execute()
                file_name = file_metadata['name']
                
                print(f"Downloading file {file_name} from Google Drive...")
                file_path = download_google_drive_file(google_drive_file_id, file_name)
                
                if file_path and os.path.exists(file_path):
                    print(f"File {file_name} saved successfully to {file_path}")
                else:
                    print(f"Failed to save file {file_name}")
            except Exception as e:
                print(f"Error retrieving file metadata for Google Drive ID {google_drive_file_id}: {e}")
        else:
            print(f"No Google Drive File ID found for record: {record_id}")

if __name__ == "__main__":
    # Replace 'YOUR_CUSTOM_OBJECT_API_NAME' with the Salesforce custom object API name
    custom_object_api_name = 'CustomerInfo__c'
    save_salesforce_files(custom_object_api_name)