"""
CircuitPack packager - create CircuitPack packages
"""

import os
import pathlib
import hashlib
from typing import Iterable, Union


def pack(contents: dict[str, bytes]) -> bytes:
    """
    Pack the data from `contents`, a `dict` with paths (as `str`s) given as
    keys and file contents (as `bytes`es) as values, and return the archive
    as a `bytes`.
    """

    header = b"CPAv001\n"

    index = b""
    data = b""
    for name, content in contents.items():
        index += f"{name} {len(content)}\n".encode("utf-8")
        data += content

    length = f"{len(index):07}\n".encode("utf-8")

    return header + length + index + data


def pack_files(
    paths: Iterable[Union[str, os.PathLike[str]]],
    directory: Union[str, os.PathLike[str]] = ".",
) -> bytes:
    """
    Pack the files specified in `paths`, relative to `directory` (which
    corresponds to the root of the device), and return the archive as a
    `bytes`.
    """

    contents = {}
    for path in paths:
        with open(os.path.join(directory, path), "rb") as file_to_pack:
            contents[str(path)] = file_to_pack.read()
    return pack(contents)


def pack_dir(directory: Union[str, os.PathLike[str]]) -> bytes:
    """
    Pack all files in `directory` (which corresponds to the root of the
    device), and return the archive as a `bytes`.
    """

    paths = []

    directory = pathlib.Path(directory)

    for root, _, files in os.walk(directory):
        for file in files:
            path = pathlib.Path(root, file).relative_to(directory)
            paths.append(path)

    return pack_files(paths, directory)


def name_file(name: str, packed: bytes) -> str:
    """
    Find and return the file name for the package named `name` with contents
    `packed`.
    """
    return name + "." + hashlib.sha256(packed).hexdigest()[:16] + ".cpa"


def pack_to_file(
    name: str,
    directory: Union[str, os.PathLike[str]],
    output: Union[str, os.PathLike[str]] = ".",
) -> str:
    """
    Pack all files in `directory` (which corresponds to the root of the
    device), find the file name based on the package name (`name`) and the hash
    of the archive, save the archive under this name in directory `output`
    (defaults to the current directory), and return the path to the archive.
    """

    packed = pack_dir(directory)

    path = os.path.join(output, name_file(name, packed))

    with open(path, "wb+") as package_file:
        package_file.write(packed)

    return path


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(prog="cppack")

    parser.add_argument("name", help="package name")
    parser.add_argument(
        "directory", help="directory to pack; corresponds to device root"
    )
    parser.add_argument(
        "--output", help="directory in which to place package", default="."
    )

    args = parser.parse_args()

    print(pack_to_file(args.name, args.directory, args.output))