from zipfile import ZipFile
import zipfile
from pathlib import Path
import base64
import sys
import io


def write_folder(zfile: ZipFile, dir_path: Path, prefix: str = ""):
    assert dir_path.is_dir()
    prefix += dir_path.name + "/"
    for sub_path in dir_path.iterdir():
        if sub_path.is_dir():
            if sub_path.name[0] in (".", "_") or sub_path.name == "tests":
                continue
            write_folder(zfile, sub_path, prefix)
        elif sub_path.name.endswith(".py"):
            zfile.write(
                sub_path,
                prefix + sub_path.name,
                zipfile.ZIP_DEFLATED)


if __name__ == "__main__":
    """Usage: python [name_of_the_script].py path_to_package1 path_to_package2 ... path_to_packageN

    The base64-encoded results will be saved to /tmp/out.txt
    Paste the results into a Kaggle Kernel as a byte object, base64 decode it, and finally unzip.
    """
    dir_paths = [Path(x) for x in sys.argv[1:]]
    assert all([x.exists() for x in dir_paths])
    zip_stream = io.BytesIO()
    zfile = ZipFile(zip_stream, 'w', compression=zipfile.ZIP_DEFLATED)
    for dir_path in dir_paths:
        write_folder(zfile, dir_path)
    print("\n".join(zfile.namelist()))
    zfile.close()
    zip_stream.seek(0)
    result = base64.encodebytes(zip_stream.getvalue())
    with open("/tmp/out.txt", "wb") as fout:
        fout.write(result)
    print("Results written to /tmp/out.txt")