#!/usr/bin/python3

from argparse import ArgumentParser
from binascii import hexlify
from functools import total_ordering
from hashlib import sha256
from importlib.util import cache_from_source
from os import chdir, walk
from pathlib import Path


def main():
    args = parse_args()
    init = "__init__.py"

    # Look for files (i.e. Python sources) that need to be unbundled.
    chdir(args.buildroot_sitelib)
    for root, dirs, files in walk(".", topdown=False):
        if Path(root).name == "__pycache__":
            continue

        removed_files = set()
        for filename in files:
            unbundled_path = args.sitelib / root / filename
            if unbundled_path.exists():
                our_path = args.buildroot_sitelib / root / filename
                with our_path.open("rb") as our_file:
                    our_hash = Checksum(our_file)
                with unbundled_path.open("rb") as unbundled_file:
                    unbundled_hash = Checksum(unbundled_file)
                if our_hash != unbundled_hash:
                    relative = Path(root) / filename
                    print(f"*** Differs from dependency copy: {relative}")
                Path(cache_from_source(our_path)).unlink()
                our_path.unlink()
                removed_files.add(filename)
                relative = Path(root) / filename
                print(f"Unbundled {relative}")
        files = set(files).difference(removed_files)

        # Do not install __init__.py in a directory that exists in a dependent
        # package.  Typically this means the directory is a namespace package,
        # and the dependent package didn’t install an __init__.py but this
        # package did.
        if init in files and (args.sitelib / root).exists():
            init_path = args.buildroot_sitelib / root / init
            Path(cache_from_source(init_path)).unlink()
            init_path.unlink()
            files.remove(init)
            print(f"Unbundled {Path(root) / init}")

        if not files and not dirs:
            # Everything in this directory was removed, so we should remove it
            # too.
            (args.buildroot_sitelib / root).rmdir()
            print(f"Completely unbundled {Path(root)}")


@total_ordering
class Checksum:
    def __init__(self, fileobj, chunk_size=256 * 1024):
        m = sha256()
        for chunk in iter(lambda: fileobj.read(chunk_size), b""):
            m.update(chunk)
        self._digest = m.digest()

    def __bytes__(self):
        return self._digest

    def __str__(self):
        return hexlify(self._digest)

    def __repr__(self):
        return f"Checksum(<{self}>)"

    def __eq__(self, other):
        return self._digest == other._digest

    def __lt__(self, other):
        return self._digest < other._digest


def parse_args():
    p = ArgumentParser(
        description="Unbundle modules and packages provided by dependencies"
    )
    p.add_argument(
        "buildroot_sitelib",
        metavar="BUILDROOT_SITELIB",
        type=Path,
        help="Sitelib directory in which (only) this package was installed",
    )
    p.add_argument(
        "sitelib",
        metavar="SITELIB",
        type=Path,
        help="Sitelib directory in which dependent packages were installed",
    )
    return p.parse_args()


if __name__ == "__main__":
    main()
