import os
import json
import sys
import requests
from concurrent.futures import ThreadPoolExecutor
import zipfile
from pydrive.auth import GoogleAuth
from pydrive.drive import GoogleDrive

# Function to create directories and save a file
def download_file(file_info, base_dir):
    try:
        file_path = file_info["path"]
        friendly_url = file_info["friendly"]

        # Create full local file path
        local_path = os.path.join(base_dir, file_path)
        os.makedirs(os.path.dirname(local_path), exist_ok=True)

        # Check if file already exists
        if os.path.exists(local_path):
            print(f"File already exists, skipping: {local_path}")
            return

        # Log the download
        print(f"Downloading: {friendly_url}")

        # Download the file
        response = requests.get(friendly_url, stream=True)
        response.raise_for_status()

        with open(local_path, "wb") as file:
            for chunk in response.iter_content(chunk_size=8192):
                file.write(chunk)

        print(f"Saved to: {local_path}")
    except Exception as e:
        print(f"Error downloading {file_info['path']}: {e}")

# Function to create a ZIP file of the directory
def create_zip(directory, zip_name):
    try:
        with zipfile.ZipFile(zip_name, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for root, dirs, files in os.walk(directory):
                for file in files:
                    file_path = os.path.join(root, file)
                    arcname = os.path.relpath(file_path, start=directory)
                    zipf.write(file_path, arcname)
        print(f"ZIP file created: {zip_name}")
    except Exception as e:
        print(f"Error creating ZIP file: {e}")

# Function to upload a file to Google Drive in a specific folder
def upload_to_google_drive(zip_file, folder_id):
    try:
        # Authenticate with Google Drive
        gauth = GoogleAuth()
        gauth.CommandLineAuth()  # Authenticate using the command line
        drive = GoogleDrive(gauth)

        # Upload the file
        file_name = os.path.basename(zip_file)
        gfile = drive.CreateFile({'title': file_name, 'parents': [{'id': folder_id}]})
        gfile.SetContentFile(zip_file)
        gfile.Upload()
        print(f"Uploaded to Google Drive folder {folder_id}: {file_name}")
    except Exception as e:
        print(f"Error uploading to Google Drive: {e}")

# Function to process the JSON file and download files
def download_files_from_json(json_file, base_dir, folder_id):
    # Load the JSON data
    try:
        with open(json_file, "r") as file:
            file_list = json.load(file)
    except Exception as e:
        print(f"Error reading JSON file: {e}")
        sys.exit(1)

    # Extract bucket name from the first file entry
    if len(file_list) == 0:
        print("The JSON file is empty.")
        sys.exit(1)

    bucket_name = file_list[0]["bucket_name"]
    bucket_dir = os.path.join(base_dir, bucket_name)

    # Create the bucket directory
    os.makedirs(bucket_dir, exist_ok=True)

    # Use ThreadPoolExecutor to download files concurrently
    with ThreadPoolExecutor(max_workers=5) as executor:
        for file_info in file_list:
            executor.submit(download_file, file_info, bucket_dir)

    # Create a ZIP file of the directory
    zip_name = f"{bucket_dir}.zip"
    create_zip(bucket_dir, zip_name)

    # Upload the ZIP file to Google Drive in the specified folder
    upload_to_google_drive(zip_name, folder_id)

if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("Usage: python3 download_and_upload_with_folder.py <json_file> <download_directory> <google_drive_folder_id>")
        sys.exit(1)

    json_file = sys.argv[1]
    download_directory = sys.argv[2]
    google_drive_folder_id = sys.argv[3]

    download_files_from_json(json_file, download_directory, google_drive_folder_id)