core/drive_utils.py

import google.generativeai as genai
# Google Drive (optional)
from google.oauth2.service_account import Credentials
from googleapiclient.discovery import build
from googleapiclient.http import MediaFileUpload, MediaIoBaseDownload
import gspread
from oauth2client.service_account import ServiceAccountCredentials

import os, io, time, re, json

#––– Authentication setup –––
GDRIVE_PARENT_FOLDER_NAME = "mtDNA-Location-Classifier"
GDRIVE_DATA_FOLDER_NAME = os.environ["GDRIVE_DATA_FOLDER_NAME"]
GCP_CREDS_DICT = json.loads(os.environ["GCP_CREDS_JSON"])  # from HF secrets
GDRIVE_CREDS = Credentials.from_service_account_info(GCP_CREDS_DICT, scopes=["https://www.googleapis.com/auth/drive"])
drive_service = build("drive", "v3", credentials=GDRIVE_CREDS)

def get_or_create_drive_folder(name, parent_id=None):
    query = f"name='{name}' and mimeType='application/vnd.google-apps.folder'"
    if parent_id:
        query += f" and '{parent_id}' in parents"
    results = drive_service.files().list(q=query, spaces='drive', fields="files(id, name)").execute()
    items = results.get("files", [])
    if items:
        return items[0]["id"]
    file_metadata = {
        "name": name,
        "mimeType": "application/vnd.google-apps.folder"
    }
    if parent_id:
        file_metadata["parents"] = [parent_id]
    file = drive_service.files().create(body=file_metadata, fields="id").execute()
    return file["id"]
# def find_drive_file(filename, parent_id):
#     """
#     Checks if a file with the given name exists inside the specified Google Drive folder.
#     Returns the file ID if found, else None.
#     """
#     query = f"'{parent_id}' in parents and name = '{filename}' and trashed = false"
#     results = drive_service.files().list(q=query, spaces='drive', fields='files(id, name)', pageSize=1).execute()
#     files = results.get('files', [])
#     if files:
#         return files[0]["id"]
#     return None

def find_drive_file(filename, parent_id):
    """
    Checks if a file with the given name exists inside the specified Google Drive folder.
    Returns the file ID if found, else None.
    """
    try:
        print(f"🔍 Searching for '{filename}' in folder: {parent_id}")
        query = f"'{parent_id}' in parents and name = '{filename}' and trashed = false"
        results = drive_service.files().list(
            q=query,
            spaces='drive',
            fields='files(id, name)',
            pageSize=1
        ).execute()
        files = results.get('files', [])
        if files:
            print(f"✅ Found file: {files[0]['name']} with ID: {files[0]['id']}")
            return files[0]["id"]
        else:
            print("⚠️ File not found.")
            return None
    except Exception as e:
        print(f"❌ Error during find_drive_file: {e}")
        return None



# def upload_file_to_drive(local_path, remote_name, folder_id):
#     file_metadata = {"name": remote_name, "parents": [folder_id]}
#     media = MediaFileUpload(local_path, resumable=True)
#     existing = drive_service.files().list(q=f"name='{remote_name}' and '{folder_id}' in parents", fields="files(id)").execute().get("files", [])
#     if existing:
#         drive_service.files().delete(fileId=existing[0]["id"]).execute()
#     file = drive_service.files().create(body=file_metadata, media_body=media, fields="id").execute()
#     result = drive_service.files().list(q=f"name='{remote_name}' and '{folder_id}' in parents", fields="files(id)").execute()
#     if not result.get("files"):
#         print(f"❌ Upload failed: File '{remote_name}' not found in folder after upload.")
#     else:
#         print(f"✅ Verified upload: {remote_name}")
#     return file["id"]
def upload_file_to_drive(local_path, remote_name, folder_id):
    try:
        if not os.path.exists(local_path):
            raise FileNotFoundError(f"❌ Local file does not exist: {local_path}")

        # Delete existing file on Drive if present
        existing = drive_service.files().list(
            q=f"name='{remote_name}' and '{folder_id}' in parents and trashed = false",
            fields="files(id)"
        ).execute().get("files", [])

        if existing:
            drive_service.files().delete(fileId=existing[0]["id"]).execute()
            print(f"🗑️ Deleted existing '{remote_name}' in Drive folder {folder_id}")

        file_metadata = {"name": remote_name, "parents": [folder_id]}
        media = MediaFileUpload(local_path, resumable=True)
        file = drive_service.files().create(
            body=file_metadata,
            media_body=media,
            fields="id"
        ).execute()

        print(f"✅ Uploaded '{remote_name}' to Google Drive folder ID: {folder_id}")
        return file["id"]

    except Exception as e:
        print(f"❌ Error during upload: {e}")
        return None


def download_file_from_drive(remote_name, folder_id, local_path):
    results = drive_service.files().list(q=f"name='{remote_name}' and '{folder_id}' in parents", fields="files(id)").execute()
    files = results.get("files", [])
    if not files:
        return False
    file_id = files[0]["id"]
    request = drive_service.files().get_media(fileId=file_id)
    fh = io.FileIO(local_path, 'wb')
    downloader = MediaIoBaseDownload(fh, request)
    done = False
    while not done:
        _, done = downloader.next_chunk()
    return True
def download_drive_file_content(file_id):
    request = drive_service.files().get_media(fileId=file_id)
    fh = io.BytesIO()
    downloader = MediaIoBaseDownload(fh, request)
    done = False
    while not done:
        _, done = downloader.next_chunk()
    fh.seek(0)
    return fh.read().decode("utf-8")