#!/usr/bin/env python3

import argparse
import collections
import os
import subprocess

import pygit2

try:
    from pygit2 import GIT_OBJECT_TREE
except ImportError:
    from pygit2.enums import ObjectType
    GIT_OBJECT_TREE = ObjectType.TREE


"""Work around not being able to commit empty directories in git-ubuntu

For details, see:

    https://canonical-git-ubuntu.readthedocs-hosted.com/en/latest/howto/restore-empty-directories.html
"""


# Modeled on gitubuntu.GitUbuntuRepository._add_missing_tree_dirs
def _add_empty_dirs(repo, tree, empty_dirs, _sub_path=()):
    """Return a tree modified to include the specified empty directories

    :param pygit2.Repository repo: the repository in which the tree resides
    :param pygit2.Tree tree: the tree object to modify
    :param dict(tuple(str): set(str)) empty_dirs: empty directories to include,
        including parent expansions, in the form returned by
        _expand_empty_dirs().
    :param tuple(str) _sub_path: internal use only during recursion. This is
        the directory we are currently recursing into, using the same tuple
        path form as the return value of _find_empty_dirs().
    :rtype: pygit2.Oid
    :returns: the id of a pygit2.Tree object representing the (possibly)
        modified tree.
    """
    # tree_builder is None if we don't need one yet, or is the replacement
    # for the tree object for this recursive call
    tree_builder = None

    for child in empty_dirs[_sub_path]:
        old_subtree = tree[child] if tree and child in tree else None
        new_subtree_id = _add_empty_dirs(
            repo=repo,
            _sub_path=_sub_path + (child,),
            tree=old_subtree,
            empty_dirs=empty_dirs,
        )
        if old_subtree and old_subtree.id == new_subtree_id:
            continue

        if tree_builder is None:
            if tree is None:
                tree_builder = repo.TreeBuilder()
            else:
                tree_builder = repo.TreeBuilder(repo.get(tree.id))

        if tree_builder.get(child):
            tree_builder.remove(child)
        tree_builder.insert(                 # (takes no kwargs)
            child,                           # name
            new_subtree_id,                  # oid
            pygit2.GIT_FILEMODE_TREE,        # attr
        )

    if tree_builder is None:
        if tree and len(tree):
            return tree.id
        else:
            return repo.TreeBuilder().write()
    else:
        return tree_builder.write()


def _find_empty_dirs(repo, tree, _sub_path=()):
    """Find the empty directories in a tree object

    :param pygit2.Repository repo: the repository the given tree is found in.
    :param pygit2.Tree tree: the tree object in which to find the empty
        directories.
    :rtype: sequence(tuple(str))
    :returns: sequence(tuple(str)): empty directories found in tree. Each empty
        directory is expressed as a tuple of strings, where each string is a
        "path component".
    """
    if not len(tree):
        yield _sub_path
    for entry in tree:
        if entry.type == GIT_OBJECT_TREE:
            yield from _find_empty_dirs(
                repo=repo,
                tree=repo.get(entry.id),
                _sub_path=_sub_path + (entry.name,),
            )


def _expand_empty_dirs(input_dirs):
    """Expand empty directories into a structure that includes all parents

    Expand a list of empty directories to also include all the directories that
    lead up to the empty directories. An expansion example:

        ["foo/bar/baz"] -> [["foo", "foo/bar", "foo/bar/baz"]]

    However, the input and outputs are represented differently. Paths are
    generally represented using the same format as the return value of
    _find_empty_dirs() and are documented there. The return value is a
    dictionary keyed by parent path, and the value is a set of subdirectory
    names.

    :param sequence(tuple(str)) input_dirs: empty directories in the form
        returned by _find_empty_dirs()
    :rtype: dict(tuple(str), str)
    :returns: expanded empty directories. Each empty directory comprises a
        parent directory and a name (like dirname(1) and basename(1)). The key
        is the parent directory, broken down into components and provided as a
        tuple. The value is the set of basenames of empty directories to create
        for the given parent directory.
    """
    result = collections.defaultdict(set)
    for input_dir in input_dirs:
        for i in range(len(input_dir)):
            result[input_dir[:i]].add(input_dir[i])
    return result


def fixup_commit(repo, commit, empty_dirs):
    assert len(commit.parents) == 1
    return repo.create_commit(                              # (takes no kwargs)
        None,                                               # ref
        commit.author,                                      # author
        commit.committer,                                   # committer
        commit.message,                                     # message
        _add_empty_dirs(
            repo=repo,
            tree=commit.tree,
            empty_dirs=empty_dirs,
        ),                                                  # tree
        [parent.id for parent in commit.parents],           # parents
    )


def main_fix_many(args):
    subprocess.check_call([
        'git',
        'rebase',
        '-x',
        f"python3 {__file__} fix-head",
        args.base_commit,
    ])


def main_fix_head(args):
    repo = pygit2.Repository(pygit2.discover_repository(os.getcwd()))
    old_commit = repo.head.peel(pygit2.Commit)
    assert len(old_commit.parents) == 1, "HEAD may not be a merge commit"
    flat_empty_dirs = list(
        _find_empty_dirs(repo, old_commit.parents[0].peel(pygit2.Tree)),
    )
    print(f"Old HEAD was {old_commit.id!s}")
    print("Reinstating the following empty directories:")
    printable_flat_empty_dirs = ('/'.join(path) for path in flat_empty_dirs)
    print("\n".join(f"    {path}" for path in printable_flat_empty_dirs))
    fixed_commit = fixup_commit(
        repo=repo,
        commit=old_commit,
        empty_dirs=_expand_empty_dirs(flat_empty_dirs),
    )
    repo.head.set_target(
        fixed_commit,
        f"Empty directory fixup of {old_commit.id!s}",
    )
    print(f"New HEAD is {fixed_commit!s}")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(required=True)

    fix_head = subparsers.add_parser('fix-head')
    fix_head.set_defaults(func=main_fix_head)

    fix_many = subparsers.add_parser('fix-many')
    fix_many.set_defaults(func=main_fix_many)
    fix_many.add_argument('base_commit')

    args = parser.parse_args()
    args.func(args)
