import py


class limelet(object):
    """A limelet is a local greenlet-like object:

        limelet(None) => builds a new main limelet.
        limelet(parent, run) => builds a limelet with the given parent limelet.

    Each independent tree of limelets works like the single global tree
    of greenlets: in each tree, there is exactly one main limelet (the
    root) and exactly one current limelet.  If you switch to a limelet,
    you change which limelet is the current one in that tree.  The other
    trees are not affected.
    """
    def __init__(self, parent, run=None):
        if parent is None:
            tree = limetree(self)
            if run is not None:
                raise TypeError("unexpected 2nd argument with parent=None")
            self._greenlet = None
        elif isinstance(parent, limelet):
            tree = parent._tree
            if run is not None:
                self.run = run
            self._greenlet = limegreenlet(self)
        else:
            raise TypeError("'parent' argument must be a limelet, "
                            "got '%s' instead" % (
                getattr(parent, '__class__', type(parent)).__name__,))
        self._parent = parent
        self._tree = tree

    def switch(self, *args):
        if self._greenlet is not None:
            while self._greenlet.dead:
                self = self._parent
        cur = self._tree.current
        assert cur._greenlet is None, "internal limelet state corrupted"
        cur._greenlet = limegreenlet.getcurrent()
        self._tree.current = self
        g = self._greenlet
        self._greenlet = None
        return g.switch(*args)

    def getcurrent(self):
        """Get the limelet that is the current one in this limelet's tree."""
        return self._tree.current

    def _setparent(self, nparent):
        if not isinstance(nparent, limelet):
            raise TypeError("parent must be a limelet")
        if nparent._tree is not self._tree:
            raise ValueError("parent must be a limelet from the same tree")
        p = nparent
        while p is not None:
            if p is self:
                raise ValueError("cyclic parent chain")
            p = p._parent
        self._parent = nparent

    def _getparent(self):
        return self._parent

    parent = property(_getparent, _setparent)

    def _getgreenlet(self):
        if self._greenlet is None:
            return limegreenlet.getcurrent()
        else:
            return self._greenlet

    def _get_fr_frame(self):
        return self._getgreenlet().gr_frame
    gr_frame = property(_get_fr_frame)

    def __nonzero__(self):
        return bool(self._getgreenlet())

    def _get_dead(self):
        return self._getgreenlet().dead
    dead = property(_get_dead)

    def throw(self, typ=None, val=None, tb=None):
        if typ is val is None:
            typ = self._tree.localexception
        return self._getgreenlet.throw(typ, val, tb)

# ____________________________________________________________
# Internal classes

class limetree(object):
    localexception = py.magic.greenlet.GreenletExit
    def __init__(self, current):
        self.current = current

class limegreenlet(py.magic.greenlet):

    def __init__(self, limelet):
        self.baselimelet = limelet

    def run(self, *args):
        try:
            try:
                runner = self.baselimelet.run
                try:
                    del self.baselimelet.run
                except AttributeError:
                    pass
                return runner(*args)
            except self.baselimelet._tree.localexception, e:
                return e
        finally:
            return_to = self.baselimelet._parent
            self.parent = return_to._greenlet
            self.baselimelet._tree.current = return_to
            return_to._greenlet = None

# ____________________________________________________________
# Some tests

def test_simple():
    main = limelet(None)
    def do_stuff(x, y):
        return x * y
    lim = limelet(main, do_stuff)
    res = lim.switch(6, 7)
    assert res == 42


def test_composable():
    main1 = limelet(None)
    def producevalue1():
        for i in range(10):
            main1.switch(i)
    producer1 = limelet(main1, producevalue1)

    res = producer1.switch()
    assert res == 0
    res = producer1.switch()
    assert res == 1

    main2 = limelet(None)
    def do_stuff():
        x = producer1.switch()
        y = producer1.switch()
        return x * 10 + y
    lim2 = limelet(main2, do_stuff)
    res = lim2.switch()
    assert res == 23


def test_return():
    main1 = limelet(None)
    def producevalue1():
        return 42
    producer1 = limelet(main1, producevalue1)

    main2 = limelet(None)
    def do_stuff():
        x = producer1.switch()
        return x * 10
    lim2 = limelet(main2, do_stuff)
    res = lim2.switch()
    assert res == 420


def test_three_limelets():
    def func1():
        sub2.switch()
        return 1
    def func2():
        return 2
    main = limelet(None)
    for i in range(5):
        sub1 = limelet(main, func1)
        sub2 = limelet(main, func2)
        res = sub1.switch()
        assert res == 2

        sub1 = limelet(main, func1)
        sub2 = limelet(sub1, func2)
        res = sub1.switch()
        assert res == 1

        sub1 = limelet(main, func1)
        sub2 = limelet(main, func2)
        assert sub2.parent is main
        sub2.parent = sub1
        res = sub1.switch()
        assert res == 1