import subprocess
import sys
import os


bucket = "tri-ml-sandbox-16011-us-west-2-datasets"
prefix = f"cv_datasets/processed"


def download(path):
    print(f'### DOWNLOADING {path}')
    os.system(f'bash custom/prepare/sync/download_cv_processed.sh {path}')


def list_s3(s3_path):
    proc = subprocess.Popen(f'aws s3 ls {s3_path}/', stdout=subprocess.PIPE, shell=True)
    (data, _) = proc.communicate()
    data = data.decode().replace('/\n', ' ').replace('\n', ' ').replace('  ', ' ').split(' ')
    out = [data[i+1] for i in range(len(data)) if data[i] == 'PRE']
    proc.kill()
    return out


args = f'{sys.argv[1]}'
nested = int(sys.argv[2]) if len(sys.argv) == 3 else 0
    

def run(f0):
    if nested == 0:
        download(f0)
    else:
        folders1 = list_s3(f's3://{bucket}/{prefix}/{f0}')
        for f1 in folders1:
            f1 = f'{f0}/{f1}'
            if nested == 1:
                download(f1)
            else:
                folders2 = list_s3(f's3://{bucket}/{prefix}/{f1}')
                for f2 in folders2:
                    f2 = f'{f1}/{f2}'
                    if nested == 2:
                        download(f2)
                    else:
                        folders3 = list_s3(f's3://{bucket}/{prefix}/{f2}')
                        for f3 in folders3:
                            f3 = f'{f2}/{f3}'
                            if nested == 3:
                                download(f3)


if '*' in args:
    base, filter = args.split('/')
    folders = list_s3(f's3://{bucket}/{prefix}/{base}')
    params = filter.split('*')
    if params[0] != '':
        folders = [f for f in folders if f.startswith(params[0])]
        params = params[1:]
    if params[-1] != '':
        folders = [f for f in folders if f.endswith(params[0])]
        params = params[:-1]
    for p in params:
        folders = [f for f in folders if p in f]
    for f in folders:
        run(f'{base}/{f}')
else:
    run(args)