from mercurial import error

from mercurial import util

from mercurial import destutil

from mercurial import extensions

from mercurial import bookmarks

from mercurial.i18n import _

def _destmergebranch(orig, repo, action='merge', sourceset=None, onheadcheck=True):

    p1 = repo['.']

    top = p1.topic()

    if top:

        heads = repo.revs('heads(topic(.)::topic(.))')

        if p1.rev() not in heads:

            raise error.Abort(_("not at topic head, update or explicit"))

        elif 1 == len(heads):

            # should look at all branch involved but... later

            bhead = ngtip(repo, p1.branch(), all=True)

            if not bhead:

                raise error.Abort(_("nothing to merge"))

            elif 1 == len(bhead):

                return bhead.first()

            else:

                raise error.Abort(_("branch '%s' has %d heads - "

                                   "please merge with an explicit rev")

                                 % (p1.branch(), len(bhead)),

                                 hint=_("run 'hg heads .' to see heads"))

        elif 2 == len(heads):

            heads = [r for r in heads if r != p1.rev()]

            # XXX: bla bla bla bla bla

            if 1 < len(heads):

                raise error.Abort(_('working directory not at a head revision'),

                                 hint=_("use 'hg update' or merge with an "

                                        "explicit revision"))

            return heads[0]

        elif 2 < len(heads):

            raise error.Abort(_("topic '%s' has %d heads - "

                                "please merge with an explicit rev")

                              % (top, len(heads)))

        else:

            assert False # that's impossible

    if getattr(orig, 'func_default', ()): # version above hg-3.7

        return orig(repo, action, sourceset, onheadcheck)

    else:

        return orig(repo)

def _destupdatetopic(repo, clean, check):

    """decide on an update destination from current topic"""

    movemark = node = None

    topic = repo.currenttopic

    revs = repo.revs('.::topic("%s")' % topic)

    if not revs:

        return None, None, None

    node = revs.last()

    if bookmarks.isactivewdirparent(repo):

        movemark = repo['.'].node()

    return node, movemark, None

def setupdest():

    if util.safehasattr(destutil, '_destmergebranch'):

        extensions.wrapfunction(destutil, '_destmergebranch', _destmergebranch)

    rebase = extensions.find('rebase')

    if (util.safehasattr(rebase, '_destrebase')

            # logic not shared with merge yet < hg-3.8

            and not util.safehasattr(rebase, '_definesets')):

        extensions.wrapfunction(rebase, '_destrebase', _destmergebranch)

    if util.safehasattr(destutil, 'destupdatesteps'):

        bridx = destutil.destupdatesteps.index('branch')

        destutil.destupdatesteps.insert(bridx, 'topic')

        destutil.destupdatestepmap['topic'] = _destupdatetopic

def ngtip(repo, branch, all=False):

    """tip new generation"""

    ## search for untopiced heads of branch

    # could be heads((::branch(x) - topic()))

    # but that is expensive

    #

    # we should write plain code instead

    subquery = '''heads(

                    parents(

                       ancestor(

                         (head() and branch(%s)

                         or (topic() and branch(%s)))))

                   ::(head() and branch(%s))

                   - topic())'''

    if not all:

        subquery = 'max(%s)' % subquery

    return repo.revs(subquery, branch, branch, branch)