import tarfile
import argparse
import os

"""make_tarfile
   Compress all files under src_dir into one tarfile output_tar
   
   Params:
     output_tar: target tar file, e.g., test.tar.gz
     src_dir: file dir to compress
"""
def make_tarfile(output_tar, src_dir):
    print(output_tar)
    print(src_dir)
    with tarfile.open(output_tar, "w:gz") as tar:
        tar.add(src_dir, arcname=os.path.basename(src_dir))
    tar.close()

"""untar_file
   Uncompress tar file compressed by make_tarfile

   Params:
     src_tar: src tar file, e.g., test.tar.gz
     output_path: target dir
"""
def untar_file(src_tar, output_path):
    print(src_tar)
    print(output_path)
    with tarfile.open(src_tar, "r:gz") as tar:
        tar.extractall(output_path)
    tar.close()


"""make_tarfile_autosplit
   Compress all files under src_dir into multiple tar files, 
     we compress 1000 files in single tarfile, it will generate 
     a list of tarfiles

   Params:
     output_dir: A dir to store all generated tarfiles
     src_dir: A dir containing all files to compress
"""
def make_tarfile_autosplit(output_dir, src_dir):
    print(output_dir)
    print(src_dir)
    files = os.listdir(src_dir)
    cnt = len(files)
    print(cnt)

    i = 0
    while True:
        if i % 1000 == 0:
            tarname = os.path.join(output_dir, (str(i)+".tar.gz"))
            print(tarname)
            with tarfile.open(tarname, "w:gz") as tar:
                for x in range (0, 1000):
                    filename = os.path.join(src_dir, files[i])
                    tar.add(filename, arcname=files[i])
                    i = i + 1
                    if i == cnt:
                        break
                tar.close()
        if i == cnt:
            break

"""untar_autosplit
   Uncompress tar files generated by make_tarfile_autosplit

   Params:
     src_tar_dir: a directory containing all tarfiles
     output_path: target dir
"""
def untar_autosplit(src_tar_dir, output_path):
    print(src_tar_dir)
    print(output_path)

    if not os.path.exists(output_path):
        os.makedirs(output_path)
    tar_files = os.listdir(src_tar_dir)

    print(tar_files)
    for tar_file in tar_files:
        print(tar_file)
        src_tar = os.path.join(src_tar_dir, tar_file)
        with tarfile.open(src_tar, "r:gz") as tar:
            tar.extractall(output_path)
        tar.close()

"""
  How to use:
    python tar_util.py --cmd tar_split --src_dir /example/dog-cat/train/dog/ --output_tar /data/test/dog-tars/
    python tar_util.py --cmd untar_split --src_tar=/data/test/dog-tars/ --output_path=/data/test/dog/ 
"""
if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument("--cmd")
    parser.add_argument("--src_dir")
    parser.add_argument("--output_tar")
    parser.add_argument("--src_tar")
    parser.add_argument("--output_path")

    args = parser.parse_args()

    cmd_type = args.cmd

    if cmd_type == "tar":
        make_tarfile(args.output_tar, args.src_dir)

    if cmd_type == "untar":
        untar_file(args.src_tar,args.output_path)

    if cmd_type == "tar_split":
        make_tarfile_autosplit(args.output_tar, args.src_dir)

    if cmd_type == "untar_split":
        untar_autosplit(args.src_tar, args.output_path)
