"""revset to select sample of repository

Hopefully this is useful to create interesting discovery cases.
"""

import collections
import random

from mercurial.i18n import _

from mercurial import (
    registrar,
    revset,
    revsetlang,
    smartset,
)

import sortedcontainers

SortedSet = sortedcontainers.SortedSet

revsetpredicate = registrar.revsetpredicate()


@revsetpredicate(b'subsetspec("<spec>")')
def subsetmarkerspec(repo, subset, x):
    """use a shorthand spec as used by search-discovery-case

    Supported format are:

    - "scratch-count-seed": not scratch(all(), count, "seed")
    - "randomantichain-seed": ::randomantichain(all(), "seed")
    - "rev-REV": "::REV"
    """
    args = revsetlang.getargs(
        x, 0, 1, _(b'subsetspec("spec") required an argument')
    )

    spec = revsetlang.getstring(args[0], _(b"spec should be a string"))
    case = spec.split(b'-')
    t = case[0]
    if t == b'scratch':
        spec_revset = b'not scratch(all(), %s, "%s")' % (case[1], case[2])
    elif t == b'randomantichain':
        spec_revset = b'::randomantichain(all(), "%s")' % case[1]
    elif t == b'rev':
        spec_revset = b'::%d' % case[1]
    else:
        assert False, spec

    selected = repo.revs(spec_revset)

    return selected & subset


@revsetpredicate(b'scratch(REVS, <count>, [seed])')
def scratch(repo, subset, x):
    """randomly remove <count> revision from the repository top

    This subset is created by recursively picking changeset starting from the
    heads. It can be summarized using the following algorithm::

        selected = set()
        for i in range(<count>):
            unselected = repo.revs("not <selected>")
            candidates = repo.revs("heads(<unselected>)")
            pick = random.choice(candidates)
            selected.add(pick)
    """
    m = _(b"scratch expects revisions, count argument and an optional seed")
    args = revsetlang.getargs(x, 2, 3, m)
    if len(args) == 2:
        x, n = args
        rand = random
    elif len(args) == 3:
        x, n, seed = args
        seed = revsetlang.getinteger(seed, _(b"seed should be a number"))
        rand = random.Random(seed)
    else:
        assert False

    n = revsetlang.getinteger(n, _(b"scratch expects a number"))

    selected = set()
    heads = SortedSet()
    children_count = collections.defaultdict(lambda: 0)
    parents = repo.changelog._uncheckedparentrevs

    baseset = revset.getset(repo, smartset.fullreposet(repo), x)
    baseset.sort()
    for r in baseset:
        heads.add(r)

        p1, p2 = parents(r)
        if p1 >= 0:
            heads.discard(p1)
            children_count[p1] += 1
        if p2 >= 0:
            heads.discard(p2)
            children_count[p2] += 1

    for h in heads:
        assert children_count[h] == 0

    selected = set()
    for x in range(n):
        if not heads:
            break
        pick = rand.choice(heads)
        heads.remove(pick)
        assert pick not in selected
        selected.add(pick)
        p1, p2 = parents(pick)
        if p1 in children_count:
            assert p1 in children_count
            children_count[p1] -= 1
            assert children_count[p1] >= 0
            if children_count[p1] == 0:
                assert p1 not in selected, (r, p1)
                heads.add(p1)
        if p2 in children_count:
            assert p2 in children_count
            children_count[p2] -= 1
            assert children_count[p2] >= 0
            if children_count[p2] == 0:
                assert p2 not in selected, (r, p2)
                heads.add(p2)

    return smartset.baseset(selected) & subset


@revsetpredicate(b'randomantichain(REVS, [seed])')
def antichain(repo, subset, x):
    """Pick a random anti-chain in the repository

    A antichain is a set of changeset where there isn't any element that is
    either a descendant or ancestors of any other element in the set. In other
    word, all the elements are independant. It can be summarized with the
    following algorithm::

    selected = set()
    unselected = repo.revs('all()')
    while unselected:
        pick = random.choice(unselected)
        selected.add(pick)
        unselected -= repo.revs('::<pick> + <pick>::')
    """

    args = revsetlang.getargs(
        x, 1, 2, _(b"randomantichain expects revisions and an optional seed")
    )
    if len(args) == 1:
        (x,) = args
        rand = random
    elif len(args) == 2:
        x, seed = args
        seed = revsetlang.getinteger(seed, _(b"seed should be a number"))
        rand = random.Random(seed)
    else:
        assert False

    cl = repo.changelog

    # We already have cheap access to the parent mapping.
    # However, we need to build a mapping of the children mapping
    parents = repo.changelog._uncheckedparentrevs
    children_map = collections.defaultdict(list)
    for r in cl:
        p1, p2 = parents(r)
        if p1 >= 0:
            children_map[p1].append(r)
        if p2 >= 0:
            children_map[p2].append(r)
    children = children_map.__getitem__

    selected = set()
    undecided = SortedSet(cl)

    while undecided:
        # while there is "undecided content", we pick a random changeset X
        # and we remove anything in `::X + X::` from undecided content
        pick = rand.choice(undecided)
        selected.add(pick)
        undecided.remove(pick)

        ancestors = set(p for p in parents(pick) if p in undecided)
        descendants = set(c for c in children(pick) if c in undecided)

        while ancestors:
            current = ancestors.pop()
            undecided.remove(current)
            for p in parents(current):
                if p in undecided:
                    ancestors.add(p)
        while descendants:
            current = descendants.pop()
            undecided.remove(current)
            for p in children(current):
                if p in undecided:
                    ancestors.add(p)

    return smartset.baseset(selected) & subset