import os
import sys
import argparse

from glob import glob
from argparse import  Namespace
from functools import partial

from anydata.converters.utils import list_s3_recursive, list_s3, run
from anydata.converters.lbm import get_sequences, process_sequence
from anydata.sync.sync_from import download_sequences
from anydata.sync.sync_utils import multi_thread, get_seqs_subset
from anydata.converters.utils import clean_string


def parse_args(num_procs=16):
    parser = argparse.ArgumentParser()
    parser.add_argument('task', type=str, nargs="+")
    parser.add_argument('--num_procs', type=int, default=num_procs)
    parser.add_argument('--verbose', action='store_true')
    parser.add_argument("--suffix", type=str, default='observations.npz')
    parser.add_argument("--src_s3_bucket", type=str, default='s3://robotics-manip-lbm')
    parser.add_argument("--dst_bucket", type=str, default='s3://tri-ml-sandbox-16011-us-west-2-datasets')
    parser.add_argument("--src_s3_path", type=str, default='efs/data/tasks')
    parser.add_argument("--dst_path_s3", type=str, default='cv_unified')
    parser.add_argument("--dst_path_local", type=str, default='/data/cv_downloaded')
    parser.add_argument("--dataset", type=str, default='LBM')
    parser.add_argument('--download_only', action='store_true')
    parser.add_argument('--overwrite', action='store_true')
    parser.add_argument('--rollout', action='store_true')

    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.dst_path_local += f'/{args.dataset}'
    args.dst_path_s3 += f'/{args.mode}/{args.dataset}'

    return args


def synchronize(args):

    source = 'teleop' if not args.rollout else 'rollout'

    # Get available tasks from the LBM spartan list
    seqs = list_s3(f'{args.src_s3_bucket}/{args.src_s3_path}/{args.task}') # Check for stations
    updated_seqs = []
    for s in seqs:
        out = list_s3(f'{args.src_s3_bucket}/{args.src_s3_path}/{args.task}/{s}')
        out = [f'{args.src_s3_bucket}/{args.src_s3_path}/{args.task}/{s}/{o}' for o in out]
        updated_seqs.extend(out)
    seqs = updated_seqs # Check for sim/real
    updated_seqs = []
    for s in seqs:
        out = list_s3_recursive(f'{s}/bc/{source}', suffix=args.suffix)
        updated_seqs.extend(out)
    seqs = updated_seqs # Check for episodes

    # Include buckets
    seqs = [f'{args.src_s3_bucket}/{os.path.dirname(s)}' for s in seqs]
    src = f'{args.src_s3_bucket}/{args.src_s3_path}'
    dst = f'{args.dst_bucket}/{args.dst_path_s3}'

    # Remove invalid paths
    seqs = [s for s in seqs if \
        ('teleop' in s or 'rollout' in s) and \
        ('duplicate' not in s and 'random' not in s)
    ]

     # If overwriting, does not check for done episodes
    if not args.overwrite:
        
        # Check for the ones already one on the unified s3
        done = list_s3_recursive(f'{dst}/{args.task}', suffix='metadata.json')
        done = [os.path.dirname(d) for d in done]

        # Check for the ones with error on the tmp folder
        # TODO @vitor inline hardcoded path may be confusing
        error = glob(f'/data/{args.dst_path_s3}/tmp/*.txt')
        error = [f'{args.dst_path_s3}/{os.path.basename(e)[:-len(".txt")]}' for e in error]
        error = ['/'.join(e.split('__')) for e in error]

        # Filter the ones already done, keeping the ones with error
        filtered_seqs = []
        for s in seqs:
            seq = s.split('/')[6:]
            seq = '/'.join([seq[i] for i in range(len(seq)) if i in [0,1,2,4,5,7]])
            seq = clean_string(f'{args.dst_path_s3}/{seq}')
            if seq in error or seq not in done: filtered_seqs.append(s)

        print('##########################')
        done, total = len(filtered_seqs), len(seqs)
        if done == 0:
            print(f'### ALL *{args.task}* EPISODES ({total}) HAVE BEEN UNIFIED ALREADY')
        else:
            print(f'### {done} REMAINING *{args.task}* EPISODES OUT OF {total} ({total - done} COMPLETED)')
        print('##########################')

        # Update sequence list
        seqs = filtered_seqs

    # Downloading parameters
    args_download = Namespace(
        task=args.task,
        src=args.src_s3_bucket,
        dst=args.dst_path_local,
        num_procs=args.num_procs,
        upload_and_delete=not args.download_only,
        verbose=args.verbose,
        mode=args.mode,
    )

    # Multi-thread download
    multi_thread_fn = partial(multi_thread, name='SYNCING LBM', 
        fn_seqs=download_sequences, fn_files=None)
    multi_thread_fn(seqs, None, args_download)


if __name__ == "__main__":

    args = parse_args()

    # Get all tasks
    tasks = []
    for task in args.task:
        if task.endswith('*'):
            options = list_s3(f'{args.src_s3_bucket}/{args.src_s3_path}')
            tasks.extend([o for o in options if o.startswith(task[:-1])])
        elif task.startswith('*'):
            options = list_s3(f'{args.src_s3_bucket}/{args.src_s3_path}')
            tasks.extend([o for o in options if o.endswith(task[1:])])
        else:
            tasks.append(task)

    # Loop over tasks
    for task in tasks:
        args.task = task
        synchronize(args)