# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import argparse
import subprocess
import numpy as np
import multiprocessing

from glob import glob
from tqdm import tqdm

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

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

def parse_args(num_procs=24):
    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('--downloaded', action='store_true')
    parser.add_argument('--downloaded_zips', action='store_true')
    parser.add_argument("--nested", type=int, default=0)
    parser.add_argument('--s3_bucket', type=str, default='s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_downloaded')
    args = parser.parse_args()

    args.src, args.dst = args.path

    if args.downloaded:
        args.dst = f'{args.s3_bucket}/{args.dst}'
    if args.downloaded_zips:
        args.dst = f'{args.s3_bucket}_zips/{args.dst}'
    args.extras = get_extras_base(args)

    return args

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

def upload_files(files, s3_path):
    for file in files:
        aws_s3_cp(file, f'{s3_path}/',
            quiet=not args.verbose,
            extras='--acl bucket-owner-full-control',
        )

def upload_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}')
        upload_sequence(i, seq, args)

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

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

if __name__ == "__main__":

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

    if len(seqs) == 0:
        print(f'########################################')
        print(f'### SYNCING SINGLE FOLDER')
        print(f'### SRC: {args.src}')
        print(f'### DST: {args.dst}')
        aws_s3_sync(args.src, args.dst,
            quiet=not args.verbose,
            extras=f'{args.extras} --acl bucket-owner-full-control',
        )
        print(f'########################################')
    else:
        multi_thread = partial(multi_thread, name='SYNCING TO',
            fn_seqs=upload_sequences, fn_files=upload_files)
        if args.nested == 0:
            multi_thread(seqs, files, args)
        else:
            for i in range(len(seqs)):
                files_i, seqs_i = get_local_files_and_folders(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_local_files_and_folders(seqs_i[j])
                        multi_thread(seqs_j, files_j, args)

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