# Copyright 2026 Toyota Research Institute.  All rights reserved.

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

from glob import glob, escape as glob_escape
from tqdm import tqdm

from anydata.utils.misc import get_local_root
from anydata.utils.read import read_json
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):
    _root = get_local_root()
    parser = argparse.ArgumentParser()
    parser.add_argument("path", type=str)
    parser.add_argument("--tmp", type=str, default=f'{_root}/tmp')
    parser.add_argument('--num_procs', type=int, default=num_procs)
    parser.add_argument("--delete", action='store_true')
    parser.add_argument("--s3_bucket", type=str, default='s3://tri-ml-sandbox-16011-us-west-2-datasets')
    parser.add_argument("--data_folder", type=str, default='cv_webbed')
    parser.add_argument("--local_bucket", type=str, default=_root)
    parser.add_argument("--src", type=str, default=None)
    parser.add_argument("--dst", type=str, default=None)
    args = parser.parse_args()

    if args.src is None:
        args.src = f'{args.local_bucket}/{args.data_folder}/{args.path}'
    if args.dst is None:
        args.dst = f'{args.s3_bucket}/{args.data_folder}/{args.path}'

    os.makedirs(args.tmp, exist_ok=True)

    return args

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

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):

    local_seq = f'{args.src}/{seq}'
    s3_seq = f'{args.dst}/{seq}'
    aws_s3_sync(local_seq, s3_seq)

    ### Delete unified
    if args.delete:
        shutil.rmtree(local_seq)

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

if __name__ == "__main__":

    args = parse_args()

    split = read_json(f'{args.src}/split_all.json')
    seqs = sorted(split['sequences'].keys())
    files = glob(f'{glob_escape(args.src)}/*.json')

    # Multi-thread uploading of sequences and files
    multi_thread = partial(multi_thread, name='UPLOADING WEBBED',
        fn_seqs=upload_sequences, tarfiles=True)
    multi_thread(seqs, files, args)

    ### Upload splits
    splits = glob(f'{glob_escape(args.src)}/*.json')
    for split in splits:
        base = split.replace(f'{args.src}/', '')
        aws_s3_cp(split, f'{args.dst}/{base}')

    ### Delete dataset folders
    if args.delete:
        shutil.rmtree(args.src)

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