import os
import pickle
from googleapiclient.discovery import build
from googleapiclient.http import MediaFileUpload
from google_auth_oauthlib.flow import InstalledAppFlow
from google.auth.transport.requests import Request

# If modifying these SCOPES, delete the file token.pickle.
SCOPES = ['https://www.googleapis.com/auth/drive.file']

def authenticate_drive():
    """Authenticate and return the Google Drive service."""
    creds = None
    if os.path.exists('token.pickle'):
        with open('token.pickle', '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:
            flow = InstalledAppFlow.from_client_secrets_file(
                'credentials.json', SCOPES)
            creds = flow.run_local_server(port=0)
        # Save the credentials for the next run
        with open('token.pickle', 'wb') as token:
            pickle.dump(creds, token)
    service = build('drive', 'v3', credentials=creds)
    return service

def get_or_create_folder(service, folder_name, parent_id=None):
    """Get the folder ID by name, or create it if it doesn't exist."""
    query = f"name='{folder_name}' and mimeType='application/vnd.google-apps.folder' and trashed=false"
    if parent_id:
        query += f" and '{parent_id}' in parents"
    results = service.files().list(q=query, spaces='drive', fields="files(id, name)").execute()
    items = results.get('files', [])
    if items:
        return items[0]['id']
    # Create folder
    file_metadata = {
        'name': folder_name,
        'mimeType': 'application/vnd.google-apps.folder'
    }
    if parent_id:
        file_metadata['parents'] = [parent_id]
    folder = service.files().create(body=file_metadata, fields='id').execute()
    return folder.get('id')

def upload_file_to_folder(service, file_path, folder_id):
    """Upload a file to the specified Google Drive folder. Returns the file's webViewLink."""
    file_name = os.path.basename(file_path)
    file_metadata = {
        'name': file_name,
        'parents': [folder_id]
    }
    media = MediaFileUpload(file_path, resumable=True)
    uploaded = service.files().create(
        body=file_metadata,
        media_body=media,
        fields='id, webViewLink'
    ).execute()
    return uploaded.get('webViewLink')

def upload_files_to_drive(file_paths, date_str=None, article_folder=None, folder_name="Farmonaut AI Videos"):
    """
    Authenticate and upload a list of files to a nested Google Drive folder structure:
    Farmonaut AI Videos/YYYY-MM-DD/article_folder/
    Returns a dict of {file_path: webViewLink}
    """
    service = authenticate_drive()
    # Step 1: Get/create the root folder
    root_id = get_or_create_folder(service, folder_name)
    # Step 2: Get/create the date folder
    if date_str:
        date_id = get_or_create_folder(service, date_str, parent_id=root_id)
    else:
        date_id = root_id
    # Step 3: Get/create the article folder
    if article_folder:
        article_id = get_or_create_folder(service, article_folder, parent_id=date_id)
    else:
        article_id = date_id
    # Step 4: Upload files to the article folder
    links = {}
    for file_path in file_paths:
        try:
            link = upload_file_to_folder(service, file_path, article_id)
            links[file_path] = link
        except Exception as e:
            links[file_path] = f"Error: {e}"
    return links

if __name__ == "__main__":
    # Example usage
    import sys
    if len(sys.argv) < 2:
        print("Usage: python gdrive_upload.py <file1> [<file2> ...]")
        exit(1)
    file_paths = sys.argv[1:]
    links = upload_files_to_drive(file_paths)
    for f, l in links.items():
        print(f"{f}: {l}")