"""
Download the UniCeption format checkpoints from the AirLab Data Server
"""

import argparse
import os

import urllib3
from minio import Minio
from minio.error import S3Error
from tqdm import tqdm


def main():
    parser = argparse.ArgumentParser(description="Download UniCeption format checkpoints from AirLab Data Server")
    parser.add_argument(
        "--folders",
        nargs="+",
        default=["encoders", "info_sharing", "prediction_heads", "examples"],
        help="List of folders to download (default: all folders). Choices: encoders, info_sharing, prediction_heads, examples",
    )
    parser.add_argument("--destination", type=str, default="./", help="Destination folder for downloaded checkpoints")
    args = parser.parse_args()

    access_key = "bT79gQYtfhpxFIitlpns"
    secret_key = "g7mSvUJ5k2a9mKv9IbhwXmUQjQX52MLwulhW9ONO"
    # client = Minio("airlab-share-02.andrew.cmu.edu:9000", access_key=access_key, secret_key=secret_key, secure=True)
    client = Minio(
        "airlab-share-02.andrew.cmu.edu:9000",
        access_key=access_key,
        secret_key=secret_key,
        secure=True,
        http_client=urllib3.PoolManager(cert_reqs="CERT_NONE"),  # disables SSL verification
    )

    bucket_name = "uniception"

    def download_folder(folder_name, bucket_name, client, destination_folder):
        folder_name = f"checkpoints/{folder_name}/"
        objects = client.list_objects(bucket_name, prefix=folder_name, recursive=True)
        for obj in tqdm(objects, desc=f"Downloading {folder_name}"):
            destination_file = os.path.join(destination_folder, obj.object_name)
            if not os.path.exists(destination_file):
                os.makedirs(os.path.dirname(destination_file), exist_ok=True)
                try:
                    client.fget_object(bucket_name, obj.object_name, destination_file)
                    print(f"Downloaded {obj.object_name} to {destination_file}")
                except S3Error as e:
                    print(f"Error downloading {obj.object_name}: {e}")
            else:
                print(f"File {destination_file} already exists. Skipping...")

    for folder in args.folders:
        download_folder(folder, bucket_name, client, args.destination)


if __name__ == "__main__":
    main()
