from __future__ import absolute_import

import functools

from mercurial import (

    discovery,

    error,

    extensions,

    phases,

    scmutil,

    util,

)

from mercurial.i18n import _

from . import (

    compat,

    exthelper,

)

eh = exthelper.exthelper()

@eh.uisetup

def uisetup(ui):

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

    extensions.wrapfunction(scmutil, 'enforcesinglehead', enforcesinglehead)

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

            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(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

def _get_branch_name(ctx):

    # make it easy for extension with the branch logic there

    branch = ctx.branch()

    if util.safehasattr(ctx, 'topic'):

        topic = ctx.topic()

        if topic:

            branch = "%s:%s" % (branch, topic)

    return branch

def _filter_obsolete_heads(repo, heads):

    """filter heads to return non-obsolete ones

    Given a list of heads (on the same named branch) return a new list of heads

    where the obsolete part have been skimmed out.

    """

    new_heads = []

    old_heads = heads[:]

    while old_heads:

        rh = old_heads.pop()

        ctx = repo[rh]

        current_name = _get_branch_name(ctx)

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

        if not ctx.obsolete():

            new_heads.append(rh)

            continue

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

        # (already filtered heads are "ignored"))

        sections_revs = repo.revs(

            b'only(%d, (%ld+%ld))', rh, old_heads, new_heads,

        )

        keep_revs = []

        for r in sections_revs:

            ctx = repo[r]

            if ctx.obsolete():

                continue

            if _get_branch_name(ctx) != current_name:

                continue

            keep_revs.append(r)

        for h in repo.revs(b'heads(%ld and (::%ld))', sections_revs, keep_revs):

            new_heads.append(h)

    new_heads.sort()

    return new_heads

def enforcesinglehead(orig, repo, tr, desc, accountclosed=False):

    """check that no named branch has multiple heads"""

    nodesummaries = scmutil.nodesummaries

    if desc in (b'strip', b'repair'):

        # skip the logic during strip

        return

    visible = repo.filtered(b'visible')

    # possible improvement: we could restrict the check to affected branch

    bm = visible.branchmap()

    cl = repo.changelog

    to_rev = cl.rev

    to_node = cl.node

    for name in bm:

        all_heads = bm.branchheads(name, closed=accountclosed)

        all_heads = [to_rev(n) for n in all_heads]

        heads = _filter_obsolete_heads(repo, all_heads)

        heads = [to_node(r) for r in heads]

        if len(heads) > 1:

            msg = _(b'rejecting multiple heads on branch "%s"')

            msg %= name

            hint = _(b'%d heads: %s')

            hint %= (len(heads), nodesummaries(repo, heads))

            raise error.Abort(msg, hint=hint)