import os
import subprocess
from tqdm import tqdm
import sys


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


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


def process(path_s3):
    folders = list_s3(path_s3)

def recursive(path_s3, n):
    # if n > 0:
    #     if 'BimanualLayCerealBoxOnCuttingBoardFromTopShelf' not in path_s3:
    #         return []
    folders = list_s3(f'{path_s3}')
    if len(folders) == 1 and folders[0] == 'tarfiles':
        pass
    elif len(folders) == 2 and folders[0] == 'indices':
        indices = list_s3(f'{path_s3}/indices')
        if 'filenames.txt' not in indices:
            all_tarfiles = []
            print(path_s3)
            for ind in tqdm(indices, ncols=192):
                name_s3 = f'{path_s3}/indices/{ind}/filenames.txt'
                name_local = name_s3.replace(f'{path_s3}/indices/', '').replace('/', '__')
                os.system(f'aws s3 cp {name_s3} {name_local} --quiet')
                with open(name_local, 'r') as file:
                    tarfiles = file.readlines()
                os.system(f'rm {name_local}')                
                all_tarfiles.extend(tarfiles)
            name_s3 = f'{path_s3}/indices/filenames.txt'
            name_local = name_s3.replace(f'{path_s3}/indices/', '').replace('/', '__')
            with open(name_local, 'w') as file:
                for t in all_tarfiles:
                    file.write(t) 
            os.system(f'aws s3 cp {name_local} {name_s3} --quiet')
            os.system(f'rm {name_local}')
            return all_tarfiles
        else:
            name_s3 = f'{path_s3}/indices/filenames.txt'
            name_local = name_s3.replace(f'{path_s3}/', '').replace('/', '__')
            os.system(f'aws s3 cp {name_s3} {name_local} --quiet')
            with open(name_local, 'r') as file:
                tarfiles = file.readlines()
            os.system(f'rm {name_local}')      
            print('++++', len(tarfiles))          
            return tarfiles
    else:
        all_tarfiles = []
        for f in folders:
            all_tarfiles.extend(recursive(f'{path_s3}/{f}', n+1))
            print('!!!', path_s3, f, len(all_tarfiles))
        return all_tarfiles
    
##########################
# path = 'DDAD/latI1_buf[40_f1]-ctx[sequence-single]_cam[CAMERA_01]_lidar_resi[192,320]'
# path = 'DL3DV/latI1_buf[40_f1]-ctx[sequence-single]_cam[0]_resi[136,240]'
# path = 'LBMv12/latIL4_buf[40_p1]-ctx[sequence-single]_cam[ALL]_resi[256,336]'
path = 'latents/LBMv12/IKPAL4_buf[40_p1]-ctx[sequence-single]_cam[ALL]_resi[-1,320,s16]'
##########################

name = f"{path.replace('/', '___')}.txt"
path_s3 = f'{s3_root}/{s3_folder}/{path}'

print()
print('###########################################')
print('######', path_s3)
print('###########################################')
print()

name = 'sim_real'
name_s3 = f'{path_s3}/{name}.txt'
name_local = name_s3.replace(f'{path_s3}/', '').replace('/', '__')

print()
print('###########################################')
print('######', name_local)
print('######', name_s3)
print('######', name)
print('###########################################')
print()

all_tarfiles = recursive(path_s3, 0)
with open(name_local, 'w') as file:
    for t in all_tarfiles:
        file.write(t) 
os.system(f'aws s3 cp {name_local} {name_s3} --quiet')
os.system(f'rm {name_local}')