import os
import cv2
from google.oauth2 import service_account
from googleapiclient.discovery import build
from googleapiclient.http import MediaIoBaseDownload, MediaFileUpload
import io
import time
from PIL import Image
from concurrent.futures import ThreadPoolExecutor
from deepface import DeepFace  # Advanced face recognition for better accuracy

# Load Google Drive API credentials
SCOPES = ['https://www.googleapis.com/auth/drive']
SERVICE_ACCOUNT_FILE = './salesforce-api-439514-d6b432a2e20e.json'  # Assuming file in current directory
credentials = service_account.Credentials.from_service_account_file(SERVICE_ACCOUNT_FILE, scopes=SCOPES)
drive_service = build('drive', 'v3', credentials=credentials)

# Folder IDs
aadhar_folder_id = '1Qtb5DYzSFE67Mbb5ZgDIqWUtdJaDD2F4'
cphotos_folder_id = '1DGeRqRbCPcfLDdEgP0h5fyX-MF8EQ8AH'
suspects_folder_id = '1N3RMhVD0OygeufLPYod6IYLtqzvlm3Jv'

def list_files_in_folder(folder_id):
    print(f"Listing files in folder ID: {folder_id}")
    query = f"'{folder_id}' in parents and trashed=false"
    try:
        results = drive_service.files().list(q=query, pageSize=1000, fields="files(id, name, mimeType, owners, parents)").execute()
        files = results.get('files', [])
        if not files:
            print(f"No files found in folder ID: {folder_id}. Ensure that the folder has files and the service account has access.")
        else:
            print(f"Found {len(files)} files in folder ID: {folder_id}.")
        return files
    except Exception as e:
        print(f"Error listing files in folder ID: {folder_id} - {e}")
        return []

def download_file(file_id, file_name):
    print(f"Attempting to download file: {file_name} (ID: {file_id})")
    if not file_name.lower().endswith(('.jpg', '.jpeg', '.png')):
        file_name += ".jpg"
    try:
        request = drive_service.files().get_media(fileId=file_id)
        file_io = io.BytesIO()
        downloader = MediaIoBaseDownload(file_io, request)
        done = False
        while not done:
            _, done = downloader.next_chunk()
        file_io.seek(0)
        with open(file_name, 'wb') as f:
            f.write(file_io.read())
        print(f"Successfully downloaded file: {file_name}")
        return file_name
    except Exception as e:
        print(f"Error downloading file ID: {file_id} - {e}")
        return None

def verify_and_fix_image(image_path):
    try:
        with Image.open(image_path) as img:
            img.verify()
        with Image.open(image_path) as img:
            img.save(image_path)
        print(f"Image verified and cleaned: {image_path}")
        return True
    except Exception as e:
        print(f"Image verification failed for {image_path}. Error: {e}")
        return False

def upload_file(file_path, folder_id):
    print(f"Uploading file: {file_path} to folder ID: {folder_id}")
    file_metadata = {'name': os.path.basename(file_path), 'parents': [folder_id]}
    media = MediaFileUpload(file_path, resumable=True)
    try:
        file = drive_service.files().create(body=file_metadata, media_body=media, fields='id').execute()
        print(f"Uploaded file ID: {file.get('id')}")
        return file.get('id')
    except Exception as e:
        print(f"Error uploading file {file_path} to folder ID: {folder_id} - {e}")
        return None

def process_images(aadhar_folder_id, cphotos_folder_id, suspects_folder_id):
    print("Starting processing of images...")
    aadhar_files = list_files_in_folder(aadhar_folder_id)
    cphotos_files = list_files_in_folder(cphotos_folder_id)

    aadhar_embeddings = []
    with ThreadPoolExecutor() as executor:
        future_to_aadhar = {executor.submit(download_file, file['id'], file['name']): file['name'] for file in aadhar_files}
        for future in future_to_aadhar:
            file_name = future_to_aadhar[future]
            try:
                file_path = future.result()
                if file_path:
                    embedding = DeepFace.represent(img_path=file_path, model_name='VGG-Face')
                    if embedding:
                        aadhar_embeddings.append((file_name, embedding))
            except Exception as e:
                print(f"Error processing Aadhar file {file_name}: {e}")

    if not aadhar_embeddings:
        print("No valid Aadhar face embeddings found. Exiting.")
        return

    unmatched_files = []
    with ThreadPoolExecutor() as executor:
        future_to_cphoto = {executor.submit(download_file, file['id'], file['name']): file for file in cphotos_files}
        for future in future_to_cphoto:
            file = future_to_cphoto[future]
            try:
                file_path = future.result()
                if not file_path:
                    continue

                embedding_cctv = DeepFace.represent(img_path=file_path, model_name='VGG-Face')
                if not embedding_cctv:
                    continue

                matched = False
                for aadhar_name, aadhar_embedding in aadhar_embeddings:
                    result = DeepFace.verify(embedding_cctv, aadhar_embedding, model_name='VGG-Face')
                    if result['verified']:
                        matched = True
                        print(f"Match found for CCTV file: {file['name']} with Aadhar file: {aadhar_name}")
                        break

                if not matched:
                    unmatched_files.append(file_path)
                    print(f"Unmatched image queued for upload: {file['name']}")
            except Exception as e:
                print(f"Error processing CCTV file {file['name']}: {e}")

    with ThreadPoolExecutor() as executor:
        for file_path in unmatched_files:
            executor.submit(upload_file, file_path, suspects_folder_id)

def main():
    start_time = time.time()
    print("Script started...")
    process_images(aadhar_folder_id, cphotos_folder_id, suspects_folder_id)
    end_time = time.time()
    print(f"Script completed in {end_time - start_time:.2f} seconds.")

if __name__ == "__main__":
    main()