

from anydata.sync.sync_utils import aws_s3_rm, remove_path

import os
from glob import glob 
from tqdm import tqdm
from threading import Thread
from anydata.sync.sync_utils import Paginator

paginator = Paginator()
max_threads = 32

# path = 'cv_webbed/videos/TriCMU/split_all__IKELA_32_str3'
# path = 'cv_webbed/videos/YAMxdof/split_all__ILA_32_str3_res-1x512'
path = 'cv_webbed/videos/DROID101/split_all3cams__IKELA_24_res-16x512'

failed = 'failed_droid.txt'

local_root = f'/data/{path}'
s3_root = f's3://tri-ml-sandbox-16011-us-west-2-datasets/{path}'

with open(failed, 'r') as f:
    data = f.readlines()
    data = [d.replace('\n', '') for d in data]
    data = [os.path.dirname(d[len(local_root)+1:]) for d in data]


def delete_data(d):
    local_root_d = f'{local_root}/{d}'
    s3_root_d = f'{s3_root}/{d}'

    files = glob(f'{local_root_d}/*')
    s3_files = paginator.ls_seq(s3_root_d)

    if len(files) > 0 or len(s3_files) > 0:
        print(files)
        print(s3_files)

    remove_path(local_root_d)
    aws_s3_rm(s3_root_d, robust=True, recursive=True)

# Prepare each snippet
threads = []
for d in tqdm(data):
    thread = Thread(target=delete_data, args=(d,))
    thread.start()
    threads.append(thread)
    # Keep number of threads in check to avoid blowing up memory
    while len(threads) > max_threads:
        threads[0].join()
        threads = threads[1:]

# Finish all threads before moving on
for thread in threads:
    thread.join()


