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

try:
    credentials = service_account.Credentials.from_service_account_file(
        SERVICE_ACCOUNT_FILE, scopes=SCOPES)
    drive_service = build('drive', 'v3', credentials=credentials)
except Exception as e:
    print(f"Error initializing Google Drive API: {e}")
    drive_service = None

# Salesforce Authentication
try:
    # Concatenate the password and security token for login
    #sf_password_with_token = 'Mahadev@9' + '6ljCATXImYQfMfRE76NZ9qpAL'  # Password + Security Token without spaces
    
    sf = Salesforce(
        username='bharatmali@sathkrutha.com',
        password= 'Mahadev@9',
        security_token = '6ljCATXImYQfMfRE76NZ9qpAL',
        domain='login'  # Use 'test' for sandbox environments
    )
    print('Salesforce Connection Successful')
except Exception as e:
    print(f"Error initializing Salesforce API: {e}")
    sf = None

# Test query to confirm object and field names
try:
    test_query = "SELECT Id FROM CustomerInfo__c LIMIT 1"
    test_records = sf.query(test_query)
    print("Test query successful, records found:", test_records['records'])
except Exception as e:
    print("Error in test query:", e)


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

def get_all_salesforce_files(custom_object_api_name):
    if sf is None:
        print("Salesforce connection is not initialized.")
        return []

    # 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"
    try:
        all_records = sf.query(query_all_records)
        if 'records' not in all_records or not all_records['records']:
            print("No records found or records list is empty.")
            return []
        return all_records['records']
    except Exception as e:
        print(f"Error querying Salesforce: {e}")
        return []


# Function to retrieve all files from the Salesforce custom object
def get_all_salesforce_files(custom_object_api_name):
    if sf is None:
        print("Salesforce connection is not initialized.")
        return []

    # 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"
    try:
        all_records = sf.query(query_all_records)
        records = all_records['records']
        return records
    except Exception as e:
        print(f"Error querying Salesforce: {e}")
        return []

# Function to download a file from Google Drive and save it to a temporary location
def download_google_drive_file(file_id, file_name):
    if drive_service is None:
        print("Google Drive service is not initialized.")
        return None

    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_Fille_ID__c']
        record_id = record['Id']
        
        # Check if the Google Drive File ID exists
        if google_drive_file_id and drive_service:
            # 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)