import os
import subprocess
from tqdm import tqdm
import sys


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(' ')
    data = [o for o in data if len(o) > 0 and o not in ['PRE']]
    proc.kill()
    return data


def list_s3_recursive(s3_path):
    proc = subprocess.Popen(f'aws s3 ls {s3_path}/ --recursive', stdout=subprocess.PIPE, shell=True)
    (data, _) = proc.communicate()
    data = data.decode().replace('/\n', ' ').replace('\n', ' ').replace('  ', ' ').split(' ')
    data = [o for o in data if len(o) > 0 and o not in ['PRE']]
    proc.kill()
    return data

s3_root = 's3://tri-ml-sandbox-16011-us-west-2-datasets'
s3_folder = 'cv_datasets/webdatasets'


def list_filenames(path, nested, filt):

    folders = list_s3_recursive(path)
    folders = sorted([f for f in folders if f.endswith('.txt')])

    if nested is not None:
        upd_folders = [] 
        for f in folders:
            split = f.split('/')
            if len(split) - 8 == nested:
                upd_folders.append(f)
        folders = upd_folders

    if filt is not None:
        upd_folders = [] 
        for f in folders:
            if filt in f:
                upd_folders.append(f)
        folders = upd_folders

    with open('gathered.txt', 'w') as file:
        for f in folders:
            file.write(f'{s3_root}/{f}\n')



if __name__ == "__main__":

    path = sys.argv[1]
    if path.endswith('/'):
        path = path[:-1]

    if len(sys.argv) > 2:
        nested = int(sys.argv[2])
    else:
        nested = 0

    if len(sys.argv) > 3:
        filt = str(sys.argv[3])
    else:
        filt = None

    list_filenames(f'{path}/indices', nested, filt)