import os
import sys
import tarfile


def build_package(package_toml_path, output_path, directories_to_include):
    data = read_package_toml(package_toml_path)
    header = build_header(data)
    pkg_path = output_path + ".pkg"

    with open(pkg_path, mode='wb') as pkg:
        pkg.write(header)
        write_manifest(pkg, data)
        build_tarball(pkg, directories_to_include)

    return pkg_path


def read_package_toml(package_toml_path):
    with open(package_toml_path, mode='r') as mani:
        return mani.read()


def build_header(data):
    header = bytes([0x01, (len(data) >> 8) & 0xFF, len(data) & 0xFF])
    return header


def write_manifest(pkg, data):
    pkg.write(data.encode("utf-8"))


def build_tarball(pkg, directories_to_include):
    with tarfile.TarFile("/tmp/pkgtar", 'w') as pkgtar:
        for directory in directories_to_include:
            add_files_to_tar(pkgtar, directory, True)

    append_tar_to_pkg(pkg)
    cleanup_tmp_files()


def add_files_to_tar(pkgtar, directory, skip_target=False):
    for root, dirs, files in os.walk(directory):
        for file in files:
            if skip_target and "target" in root:
                continue
            print(f"\33[2Kadd: {os.path.join(root, file)}", end="\r", flush=True)
            pkgtar.add(os.path.join(root, file))
    print("\33[2K", end="\r", flush=True)


def append_tar_to_pkg(pkg):
    with open("/tmp/pkgtar", 'rb') as pkgtar:
        pkg.write(pkgtar.read())


def cleanup_tmp_files():
    print("deleting /tmp/pkgtar...")
    os.unlink("/tmp/pkgtar")


if __name__ == '__main__':
    if len(sys.argv) < 3:
        print("Usage: pkg.py <package_toml_path> <output_path> [<directories_to_include> ...]")
        sys.exit(1)

    package_toml_path = sys.argv[1]
    output_path = sys.argv[2]
    directories_to_include = sys.argv[3:] if len(sys.argv) > 3 else None

    try:
        pkg_path = build_package(package_toml_path, output_path, directories_to_include)
        print(f"Package created: {pkg_path}")
    except FileNotFoundError:
        print("Error: File not found.")
    except Exception as e:
        print(f"Error occurred: {e}")