import os
import zipfile
import tarfile
import argparse

from glob import glob
from tqdm import tqdm

from anydata.sync.sync_utils import multi_thread
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('--zip', action='store_true')
    parser.add_argument('--tar', action='store_true')
    parser.add_argument('--targz', action='store_true')
    args = parser.parse_args()

    args.src, args.dst = args.path

    return args

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

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

def extract_sequence(i, seq, args):
    seq_out = os.path.splitext(seq.replace(args.src, args.dst))[0] + '/'
    os.makedirs(os.path.dirname(seq_out), exist_ok=True)
    try:
        if seq.endswith('.zip'):   
            with zipfile.ZipFile(seq, 'r') as zip_ref:
                zip_ref.extractall(path=seq_out)    
        elif seq.endswith('.tar'):
            with tarfile.open(seq, 'r') as tar_ref:
                tar_ref.extractall(path=seq_out)
    except Exception as e:
        tmp = f'{args.dst}/tmp/{seq_out.replace("/","__")}.txt'
        os.makedirs(os.path.dirname(tmp), exist_ok=True)
        with open(tmp, 'w') as fp: pass

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

if __name__ == "__main__":

    args = parse_args()
    if args.zip:
        seqs = glob(f'{args.src}/**/*.zip', recursive=True)
    elif args.tar:
        seqs = glob(f'{args.src}/**/*.tar', recursive=True)
    elif args.targz:
        seqs = glob(f'{args.src}/**/*.tar.gz', recursive=True)
    else:
        raise ValueError('Invalid extension for extraction')
    files = None

    multi_thread = partial(multi_thread, name='EXTRACTING', 
        fn_seqs=extract_sequences)
    multi_thread(seqs, files, args)

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