"""
Mercurial hook to check that all the changesets in a 3.x branch are merged
with either another cset in the same branch or a cset in the following 3.x
branch.

To use the changeset hook in a local repository, include something like the
following in its hgrc file.

[hooks]
pretxnchangegroup.checkheads = python:/home/hg/repos/hooks/checkmerge.py:hook
"""

from collections import defaultdict

from mercurial.node import bin, nullrev


def hook(ui, repo, node, **kwargs):
    changelog = repo.changelog
    # rev number of the first new cset
    start = changelog.rev(bin(node))
    end = len(changelog)
    # The rev numbers in this changegroup
    newcsets = range(start, end)
    # The rev numbers of the changegroup parents (not in the changegroup)
    parents = defaultdict(set)
    csets = defaultdict(set)

    for n in newcsets:
        cset = repo[n]
        cset_branch = cset.branch()
        # 2.x changesets are not merged, to we don't have to check them
        if cset_branch.startswith('2'):
            continue
        # take all the parents (even if technically we don't need the ones
        # from the oldest 3.x branch)
        for p in changelog.parentrevs(n):
            if p == nullrev:
                continue
            parents[cset_branch].add(repo[p])
        # take all the changesets in the 3.x branches (default is not needed)
        if cset_branch.startswith('3'):
            csets[cset_branch].add(cset)

    # get a list of all the branches
    all_branches = set(repo.branchmap())

    # This loop checks that every csid in 3.x is a parent of either another 3.x
    # cset or of a 3.x+1 (or default) cset.
    # If it's not a parent, it means that the cset has not been merged.
    needs_merge = False
    for branch in sorted(csets):
        current_csets = csets[branch]
        current_parents = parents[branch]
        next_branch = str(float(branch)+0.1)
        next_branch = next_branch if next_branch in all_branches else 'default'
        next_parents = parents[next_branch]
        for cset in current_csets:
            if cset not in next_parents and cset not in current_parents:
                needs_merge = True  # abort the transaction and rollback
                ui.warn('* %s (in %r) should be merged with %r!\n' %
                        (cset, branch, next_branch))
    return needs_merge