from __future__ import absolute_import

import functools

from mercurial import (

    discovery,

    extensions,

    phases,

)

from . import (

    compat,

    exthelper,

)

eh = exthelper.exthelper()

@eh.uisetup

def uisetup(ui):

    extensions.wrapfunction(discovery, '_postprocessobsolete', _postprocessobsolete)

def branchinfo(pushop, repo, node):

    return repo[node].branch()

# taken from 7d5455b988ec + branchinfo abstraction.

def _postprocessobsolete(orig, pushop, futurecommon, candidate_newhs):

    """post process the list of new heads with obsolescence information

    Exists as a sub-function to contain the complexity and allow extensions to

    experiment with smarter logic.

    Returns (newheads, discarded_heads) tuple

    """

    pushingmarkerfor = discovery.pushingmarkerfor

    # known issue

    #

    # * We "silently" skip processing on all changeset unknown locally

    #

    # * if <nh> is public on the remote, it won't be affected by obsolete

    #     marker and a new is created

    # define various utilities and containers

    repo = pushop.repo

    unfi = repo.unfiltered()

    torev = compat.getgetrev(unfi.changelog)

    public = phases.public

    getphase = unfi._phasecache.phase

    ispublic = lambda r: getphase(unfi, r) == public

    ispushed = lambda n: torev(n) in futurecommon

    hasoutmarker = functools.partial(pushingmarkerfor, unfi.obsstore, ispushed)

    successorsmarkers = unfi.obsstore.successors

    newhs = set()  # final set of new heads

    discarded = set()  # new head of fully replaced branch

    localcandidate = set()  # candidate heads known locally

    unknownheads = set()  # candidate heads unknown locally

    for h in candidate_newhs:

        if h in unfi:

            localcandidate.add(h)

        else:

            if successorsmarkers.get(h) is not None:

                msg = (

                    b'checkheads: remote head unknown locally has'

                    b' local marker: %s\n'

                )

                repo.ui.debug(msg % hex(h))

            unknownheads.add(h)

    # fast path the simple case

    if len(localcandidate) == 1:

        return unknownheads | set(candidate_newhs), set()

    # actually process branch replacement

    while localcandidate:

        nh = localcandidate.pop()

        current_branch = branchinfo(pushop, unfi, nh)

        # run this check early to skip the evaluation of the whole branch

        if torev(nh) in futurecommon or ispublic(torev(nh)):

            newhs.add(nh)

            continue

        # Get all revs/nodes on the branch exclusive to this head

        # (already filtered heads are "ignored"))

        branchrevs = unfi.revs(

            b'only(%n, (%ln+%ln))', nh, localcandidate, newhs

        )

        branchnodes = []

        for r in branchrevs:

            ctx = unfi[r]

            if ctx.branch() == current_branch:

                branchnodes.append(ctx.node())

        # The branch won't be hidden on the remote if

        # * any part of it is public,

        # * any part of it is considered part of the result by previous logic,

        # * if we have no markers to push to obsolete it.

        if (

            any(ispublic(torev(n)) for n in branchnodes)

            or any(torev(n) in futurecommon for n in branchnodes)

            or any(not hasoutmarker(n) for n in branchnodes)

        ):

            newhs.add(nh)

        else:

            # note: there is a corner case if there is a merge in the branch.

            # we might end up with -more- heads.  However, these heads are not

            # "added" by the push, but more by the "removal" on the remote so I

            # think is a okay to ignore them,

            discarded.add(nh)

    newhs |= unknownheads

    return newhs, discarded