def get_firefoxtrees(repo):
    for tag, node in sorted(repo.tags().items()):
        result = resolve_trees_to_uris([tag])[0]
        if not result[1]:
            continue
        tree, uri = result
        yield tag, node, tree, uri

def writefirefoxtrees(repo):
    """Write the firefoxtrees node mapping to the filesystem."""
    lines = []
    trees = {}
    for tree, node in sorted(repo.firefoxtrees.items()):
        assert len(node) == 20
        lines.append('%s %s' % (tree, hex(node)))
        trees[tree] = hex(node)

    _firefoxtreesrepo(repo).vfs.write('firefoxtrees', '\n'.join(lines))

    # Old versions of firefoxtrees stored labels in the localtags file. Since
    # this file is read by Mercurial and has no relevance to us any more, we
    # prune relevant entries from this file so the data isn't redundant with
    # what we now write.
    localtags = repo.opener.tryread('localtags')
    havedata = len(localtags) > 0
    taglines  = []
    for line in localtags.splitlines():
        line = line.strip()
        node, tag = line.split()
        tree, uri = resolve_trees_to_uris([tag])[0]
        if not uri:
            taglines.append(line)

    if havedata:
        repo.vfs.write('localtags', '\n'.join(taglines))

def treeherder(ui, repo, tree=None, rev=None, **opts):
    """Open Treeherder showing build status for the specified revision.

    The command receives a tree name and a revision to query. The tree is
    required because a revision/changeset may existing in multiple
    repositories.
    """
    if not tree:
        raise util.Abort('A tree must be specified.')

    if not rev:
        raise util.Abort('A revision must be specified.')

    tree, repo_url = resolve_trees_to_uris([tree])[0]
    if not repo_url:
        raise util.Abort("Don't know about tree: %s" % tree)

    r = MercurialRepository(repo_url)
    node = repo[rev].hex()
    push = r.push_info_for_changeset(node)

    if not push:
        raise util.Abort("Could not find push info for changeset %s" % node)

    push_node = push.last_node

    url = treeherder_url(tree, push_node[0:12])

    import webbrowser
    webbrowser.get('firefox').open(url)

def peerorrepo(orig, ui, path, *args, **kwargs):
    try:
        return orig(ui, path, *args, **kwargs)
    except RepoError:
        tree, uri = resolve_trees_to_uris([path])[0]
        if not uri:
            raise

        return orig(ui, uri, *args, **kwargs)

def pullexpand(orig, ui, repo, source='default', **opts):
    """Wraps built-in pull command to expand aliases to multiple sources."""
    for tree, uri in resolve_trees_to_uris([source]):
        result = orig(ui, repo, uri or tree, **opts)

        if result:
            return result

    return 0

def firefoxtrees(repo, proto):
    lines = []

    for tag, node in sorted(repo.tags().items()):
        if not resolve_trees_to_uris([tag])[0][1]:
            continue

        lines.append('%s %s' % (tag, hex(node)))

    return '\n'.join(lines)

def get_firefoxtrees(repo):
    """Generator for Firefox tree labels defined in this repository.

    Returns a tuple of (tag, node, tree, uri)
    """
    for tag, node in sorted(repo.firefoxtrees.items()):
        result = resolve_trees_to_uris([tag])[0]
        if not result[1]:
            continue
        tree, uri = result
        yield tag, node, tree, uri

def pushcommand(orig, ui, repo, dest=None, **opts):
    """Wraps commands.push to resolve names to tree URLs.

    Ideally we'd patch ``ui.expandpath()``. However, It isn't easy to tell
    from that API whether we should be giving out HTTP or SSH URLs.
    This was proposed and rejected as a core feature to Mercurial.
    http://www.selenic.com/pipermail/mercurial-devel/2014-September/062052.html
    """
    if isfirefoxrepo(repo):
        tree, uri = resolve_trees_to_uris([dest], write_access=True)[0]
        if uri:
            dest = uri

    return orig(ui, repo, dest=dest, **opts)

def firefoxtrees(repo, proto):
    lines = []

    if repo.ui.configbool('firefoxtree', 'servetagsfrombookmarks'):
        for name, hnode in sorted(bookmarks.listbookmarks(repo).items()):
            tree, uri = resolve_trees_to_uris([name])[0]
            if not uri:
                continue

            lines.append('%s %s' % (tree.encode('ascii'), hnode))
    else:
        for tag, node, tree, uri in get_firefoxtrees(repo):
            lines.append('%s %s' % (tag, hex(node)))

    return '\n'.join(lines)

def outgoingcommand(orig, ui, repo, dest=None, **opts):
    """Wraps command.outgoing to limit considered nodes.

    We wrap commands.outgoing rather than hg._outgoing because the latter is a
    low-level API used by discovery. Manipulating it could lead to unintended
    consequences.
    """
    tree, uri = resolve_trees_to_uris([dest])[0]
    rev = opts.get('rev')
    if uri and not rev:
        ui.status(_('no revisions specified; '
            'using . to avoid inspecting multiple heads\n'))
        opts['rev'] = '.'

    return orig(ui, repo, dest=dest, **opts)

def wrappedpullbookmarks(orig, pullop):
    """Wraps exchange._pullbookmarks.

    We remove remote bookmarks that match firefox tree tags when pulling
    from a repo that advertises the firefox tree tags in its own namespace.

    This is meant for the special unified repo that advertises heads as
    bookmarks. By filtering out the bookmarks to clients running this extension,
    they'll never pull down the bookmarks version of the tags.
    """
    repo = pullop.repo

    if isfirefoxrepo(repo) and pullop.remote.capable('firefoxtrees'):
        pullop.remotebookmarks = {k: v for k, v in pullop.remotebookmarks.items()
                                  if not resolve_trees_to_uris([k])[0][1]}

    return orig(pullop)

def revset_pushhead(repo, subset, x):
    """``pushhead([TREE])``
    Changesets that are push heads.

    A push head is a changeset that was a head when it was pushed to a
    repository. In other words, the automation infrastructure likely
    kicked off a build using this changeset.

    If an argument is given, we limit ourselves to pushes on the specified
    tree.

    If no argument is given, we return all push heads for all trees. Note that
    a changeset can be a push head multiple times. This function doesn't care
    where the push was made if no argument was given.
    """
    # We have separate code paths because the single tree path uses a single
    # query and is faster.
    if x:
        tree = revset.getstring(x, _('pushhead() requires a string argument.'))
        tree, uri = resolve_trees_to_uris([tree])[0]

        if not uri:
            raise util.Abort(_("Don't know about tree: %s") % tree)

        def pushheads():
            for push_id, head_changeset in repo.changetracker.tree_push_head_changesets(tree):
                try:
                    head = repo[head_changeset].rev()
                    yield head
                except error.RepoLookupError:
                    # There are some malformed pushes.  Ignore them.
                    continue

        # Push heads are returned in order of ascending push ID, which
        # corresponds to ascending commit order in hg.
        return subset & revset.generatorset(pushheads(), iterasc=True)
    else:
        def is_pushhead(r):
            node = repo[r].node()
            for push in repo.changetracker.pushes_for_changeset(node):
                if str(push[4]) == node:
                    return True
            return False

        return subset.filter(is_pushhead)

def fxheads(ui, repo, **opts):
    """Show last known head commits for pulled Firefox trees.

    The displayed list may be out of date. Pull before running to ensure
    data is current.
    """
    if not isfirefoxrepo(repo):
        raise util.Abort(_('fxheads is only available on Firefox repos'))

    displayer = cmdutil.show_changeset(ui, repo, opts)
    for tag, node in sorted(repo.tags().items()):
        if not resolve_trees_to_uris([tag])[0][1]:
            continue

        ctx = repo[node]
        displayer.show(ctx)

    displayer.close()

def revset_tree(repo, subset, x):
    """``tree(X)``
    Changesets currently in the specified Mozilla tree.

    A tree is the name of a repository. e.g. ``central``.
    """
    err = _('tree() requires a string argument.')
    tree = revset.getstring(x, err)

    tree, uri = resolve_trees_to_uris([tree])[0]
    if not uri:
        raise util.Abort(_("Don't know about tree: %s") % tree)

    ref = '%s/default' % tree

    head = repo[ref].rev()
    ancestors = set(repo.changelog.ancestors([head], inclusive=True))

    return [r for r in subset if r in ancestors]

def pushcommand(orig, ui, repo, dest=None, **opts):
    """Wraps commands.push to resolve names to tree URLs.

    Ideally we'd patch ``ui.expandpath()``. However, It isn't easy to tell
    from that API whether we should be giving out HTTP or SSH URLs.
    This was proposed and rejected as a core feature to Mercurial.
    http://www.selenic.com/pipermail/mercurial-devel/2014-September/062052.html
    """
    if isfirefoxrepo(repo):
        # Automatically define "review" unless it is already defined.
        if dest == 'review':
            if not ui.config('paths', 'review', None):
                dest = 'ssh://reviewboard-hg.mozilla.org/gecko'
        else:
            tree, uri = resolve_trees_to_uris([dest], write_access=True)[0]
            if uri:
                dest = uri

    return orig(ui, repo, dest=dest, **opts)

def revset_firstpushtree(repo, subset, x):
    """``firstpushtree(X)``
    Changesets that were initially pushed to tree X.
    """
    tree = revset.getstring(x, _('firstpushtree() requires a string argument.'))

    tree, uri = resolve_trees_to_uris([tree])[0]
    if not uri:
        raise util.Abort(_("Don't know about tree: %s") % tree)

    def fltr(x):
        pushes = list(repo.changetracker.pushes_for_changeset(
            repo[x].node()))

        if not pushes:
            return False

        return pushes[0][0] == tree

    return subset.filter(fltr)

def revset_pushhead(repo, subset, x):
    """``pushhead([TREE])``
    Changesets that are push heads.

    A push head is a changeset that was a head when it was pushed to a
    repository. In other words, the automation infrastructure likely
    kicked off a build using this changeset.

    If an argument is given, we limit ourselves to pushes on the specified
    tree.

    If no argument is given, we return all push heads for all trees. Note that
    a changeset can be a push head multiple times. This function doesn't care
    where the push was made if no argument was given.
    """
    # We have separate code paths because the single tree path uses a single
    # query and is faster.
    if x:
        tree = revset.getstring(x, _('pushhead() requires a string argument.'))
        tree, uri = resolve_trees_to_uris([tree])[0]

        if not uri:
            raise util.Abort(_("Don't know about tree: %s") % tree)

        heads = set(repo[r[4]].rev() for r in
            repo.changetracker.tree_push_heads(tree))

        return [r for r in subset if r in heads]

    revs = []

    for r in subset:
        node = repo[r].node()

        for push in repo.changetracker.pushes_for_changeset(node):
            if str(push[4]) == node:
                revs.append(r)
                break

    return revs

def pullcommand(orig, ui, repo, source='default', **opts):
    """Wraps built-in pull command to expand special aliases."""
    if not isfirefoxrepo(repo):
        return orig(ui, repo, source=source, **opts)

    # The special source "fxtrees" will pull all trees we've pulled before.
    if source == 'fxtrees':
        for tag, node, tree, uri in get_firefoxtrees(repo):
            res = orig(ui, repo, source=tree, **opts)
            if res:
                return res

        return 0
    elif source in MULTI_TREE_ALIASES:
        for tree, uri in resolve_trees_to_uris([source]):
            res = orig(ui, repo, source=tree, **opts)
            if res:
                return res

        return 0

    return orig(ui, repo, source=source, **opts)

def revset_firstpushtree(repo, subset, x):
    """``firstpushtree(X)``
    Changesets that were initially pushed to tree X.
    """
    tree = revset.getstring(x, _('firstpushtree() requires a string argument.'))

    tree, uri = resolve_trees_to_uris([tree])[0]
    if not uri:
        raise util.Abort(_("Don't know about tree: %s") % tree)

    revs = []

    for rev in subset:
        pushes = list(repo.changetracker.pushes_for_changeset(
            repo[rev].node()))

        if not pushes:
            continue

        if pushes[0][0] == tree:
            revs.append(rev)

    return revs

def pushtree(ui, repo, tree=None, rev=None, **opts):
    """Push changesets to a Mozilla repository.

    If only the tree argument is defined, we will attempt to push the current
    tip to the repository specified. This may fail due to pushed mq patches,
    local changes, etc. Please note we only attempt to push the current tip and
    it's ancestors, not all changesets not in the remote repository. This is
    different from the default behavior of |hg push| and is the distinguishing
    difference from that command.

    If you would like to push a non-active head, specify it with -r REV. For
    example, if you are currently on mozilla-central but wish to push inbound
    to mozilla-inbound, run `hg pushtree -r inbound/default inbound`.
    """
    if not tree:
        raise util.Abort(_('A tree must be specified.'))

    tree, uri = resolve_trees_to_uris([tree], write_access=True)[0]

    if not uri:
        raise util.Abort("Don't know about tree: %s" % tree)

    return push(ui, repo, rev=[rev], dest=uri)