# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import shutil
import argparse
import multiprocessing

from glob import glob
from tqdm import tqdm
from functools import partial

from anydata.converters.utils import get_splits, loop_over
from anydata.sync.sync_utils import multi_thread, aws_s3_cp, aws_s3_exists, create_tar
from anydata.utils.misc import get_local_root
from rich.console import Console

console = Console()

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

def parse_args():
    _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=16)
    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_unified')
    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)
    parser.add_argument("--verbose", action='store_true')
    parser.add_argument("--skip_existing", action='store_true',
                        help='Skip episodes already present on S3 (idempotent resume)')
    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):
    n_uploaded = 0
    n_skipped = 0
    progress = tqdm(seqs, ncols=96, leave=False)
    for seq in progress:
        result = upload_sequence(i, seq, args)
        if result == 'skipped':
            n_skipped += 1
            if args.verbose:
                console.print(f'[dim]  [Thread {i+1}] SKIP  {os.path.dirname(seq).replace(args.src + "/", "")}[/dim]')
        else:
            n_uploaded += 1
            if args.verbose:
                console.print(f'[green]  [Thread {i+1}] DONE  {os.path.dirname(seq).replace(args.src + "/", "")}[/green]')
        progress.set_description(f'### Thread {i+1}/{args.num_procs}  uploaded={n_uploaded} skipped={n_skipped}')

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

def upload_sequence(i, seq, args):

    ### Prepare names
    dirname = os.path.dirname(seq)
    base = dirname.replace(f'{args.src}/', '')
    name = base.replace('/','__')

    ### Skip if already fully uploaded
    if args.skip_existing and aws_s3_exists(f'{args.dst}/{base}/metadata.json'):
        return 'skipped'

    dirname_folders = glob(f'{dirname}/*')
    dirname_folders = [f for f in dirname_folders if not f.endswith('.json')]

    os.makedirs(args.tmp, exist_ok=True)
    for dirname_folder in dirname_folders:  # NOTE(bvh): typically modalities (rgb, lowdim, etc.)
        label = os.path.basename(dirname_folder)
        name_folder = f'{name}_{label}.tar.gz'
        ### Create tarfile
        create_tar(f'{args.tmp}/{name_folder}', dirname_folder)
        ### Upload tarfile to s3
        aws_s3_cp(f'{args.tmp}/{name_folder}', f'{args.dst}/{base}/{label}.tar.gz')
        ### Delete tarfile
        os.remove(f'{args.tmp}/{name_folder}')
    
    ### Upload local per-episode metadata
    # NOTE(bvh): acts as completion marker for skip logic
    aws_s3_cp(f'{dirname}/metadata.json', f'{args.dst}/{base}/metadata.json')

    ### Delete unified
    if args.delete:
        shutil.rmtree(dirname)
    return 'done'

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

if __name__ == '__main__':

    # Parse arguments
    args = parse_args()

    # Get sequences and files
    seqs = loop_over(args.src)
    files = glob(f'{args.src}/*.json')

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

    ### Upload splits and global metadata
    splits = glob(f'{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)

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