# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import shutil
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=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("--nested", type=int, default=0)
    parser.add_argument('--s3_bucket', type=str, default='s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_downloaded')
    parser.add_argument("--seqs", type=str, nargs='+', default=None)
    parser.add_argument("--subset", type=str, default=None)
    parser.add_argument('--upload_and_delete', action='store_true')
    parser.add_argument('--mode', type=str, choices=['frames', 'videos'], default=None)

    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument('--frames', action='store_const', dest='mode', const='frames')
    mode.add_argument('--videos', action='store_const', dest='mode', const='videos')
    
    args = parser.parse_args()

    args.src, args.dst = args.path
    if args.downloaded:
        args.src = f'{args.s3_bucket}/{args.src}'
    args.extras = get_extras_base(args)

    return args

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

def download_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',
            robust=True,
        )

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

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

    ### Process downloaded data (assume it's LBM for now)
    if args.upload_and_delete:
        # TODO @vitor avoid hardcoded /data paths below

        try:

            from anydata.converters.utils import parse_args as parse_args_converter, post_process
            from anydata.converters.lbm import process_sequence, parse_sequence

            args_process = parse_args_converter('lbm')
            args_process.src = '/data/cv_downloaded/LBM'
            args_process.dst = f'/data/cv_unified/{args.mode}/LBM'
            args_process.mode = args.mode
            args_process.upload_and_delete = True
            args_process.overwrite = False

            dst, tmp, reason = parse_sequence(s3_path, args_process)
            if reason == 'done': return # Skip if has been processed already
            dst = process_sequence(i, s3_path, dst, args_process)
            post_process(dst, tmp, args_process)

            ### Delete downloaded and unified data
            shutil.rmtree(s3_path, ignore_errors=True)
            shutil.rmtree(dst, ignore_errors=True)

        except Exception as e:
            if args.verbose:
                print(f'### ERROR PROCESSING {seq}: {e}')
            else:
                pass

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

if __name__ == "__main__":

    args = parse_args()
    if args.seqs is None:
        files, seqs = get_s3_files_and_folders(args.src)
    else:
        files, seqs = None, args.seqs

    if args.subset is not None:
        assert '/' in args.subset, f'Invalid subset format {args.subset}'
        skip, start = args.subset.split('/')
        seqs = seqs[int(start)::int(skip)]
        assert len(seqs) > 0, f'Invalid subset format {args.subset}'

    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',
            robust=True,
        )
        print(f'########################################')
    else:
        multi_thread = partial(multi_thread, name='SYNCING FROM',
            fn_seqs=download_sequences, fn_files=download_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(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(seqs_i[j])
                        multi_thread(seqs_j, files_j, args)

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