# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import argparse
import numpy as np

from tqdm import tqdm
from anydata.sync.sync_utils import get_extras_base, get_s3_files_and_folders, multi_thread, aws_s3_cp, aws_s3_sync
from functools import partial

####################################

def parse_args(num_procs=8):
    parser = argparse.ArgumentParser()
    parser.add_argument('path', type=str, nargs='+')
    parser.add_argument('--num_procs', type=int, default=num_procs)
    parser.add_argument('--verbose', action='store_true')
    parser.add_argument('--dryrun', action='store_true')
    parser.add_argument("--nested", type=int, default=0)
    args = parser.parse_args()

    args.src, args.dst = args.path
    args.extras = get_extras_base(args)

    return args

################################

def copy_files(files, s3_path):
    for file in files:
        aws_s3_cp(file, f'{s3_path}/', quiet=True, extras='--acl bucket-owner-full-control')

def copy_sequences(i, seqs, args):
    progress = tqdm(seqs, ncols=96, leave=False)
    for seq in progress:
        progress.set_description(f'### Thread {i+1}/{args.num_procs}')
        copy_sequence(i, seq, args)

def copy_sequence(i, seq, args):
    s3_path = f'{args.dst}/{seq.replace(args.src + "/", "")}'
    aws_s3_sync(seq, s3_path, quiet=False, extras=f'{args.extras} --acl bucket-owner-full-control')

################################

if __name__ == "__main__":

    args = parse_args()
    files, seqs = get_s3_files_and_folders(args.src)

    multi_thread = partial(multi_thread, name='SYNCING ACROSS',
        fn_seqs=copy_sequences, fn_files=copy_files)

    if args.nested == 0:
        multi_thread(seqs, files, args)
    else:
        for i in range(len(seqs)):
            files_i, seqs_i = get_s3_files_and_folders(f'{seqs[i]}')
            if args.nested == 1:
                multi_thread(seqs_i, files_i, args)
            else:
                for j in range(len(seqs_i)):
                    files_j, seqs_j = get_s3_files_and_folders(f'{seqs_i[j]}')
                    multi_thread(seqs_j, files_j, args)

################################