# Copyright (c) 2009-2010 Denis Bilenko. See LICENSE for details.

from collections import deque
from gevent.hub import GreenletExit, getcurrent
from gevent.greenlet import joinall, Greenlet
from gevent.timeout import Timeout

__all__ = ['GreenletSet', 'Pool']


class GreenletSet(object):
    """Maintain a set of greenlets that are still running.

    Links to each item and removes it upon notification.
    """
    greenlet_class = Greenlet

    def __init__(self, *args):
        assert len(args)<=1, args
        self.greenlets = set(*args)
        if args:
            for greenlet in args[0]:
                greenlet.rawlink(self.discard)
        # each item we kill we place in dying, to avoid killing the same greenlet twice
        self.dying = set()

    def __repr__(self):
        try:
            classname = self.__class__.__name__
        except AttributeError:
            classname = 'GreenletSet' # XXX check if 2.4 really uses this line
        return '<%s at %s %s>' % (classname, hex(id(self)), self.greenlets)

    def __len__(self):
        return len(self.greenlets)

    def __contains__(self, item):
        return item in self.greenlets

    def __iter__(self):
        return iter(self.greenlets)

    def add(self, greenlet):
        greenlet.rawlink(self.discard)
        self.greenlets.add(greenlet)

    def discard(self, greenlet):
        self.greenlets.discard(greenlet)
        self.dying.discard(greenlet)

    def spawn(self, *args, **kwargs):
        add = self.add
        greenlet = self.greenlet_class.spawn(*args, **kwargs)
        add(greenlet)
        return greenlet

    def spawn_link(self, *args, **kwargs):
        greenlet = self.spawn(*args, **kwargs)
        greenlet.link()
        return greenlet

    def spawn_link_value(self, *args, **kwargs):
        greenlet = self.spawn(*args, **kwargs)
        greenlet.link_value()
        return greenlet

    def spawn_link_exception(self, *args, **kwargs):
        greenlet = self.spawn(*args, **kwargs)
        greenlet.link_exception()
        return greenlet

#     def close(self):
#         """Prevents any more tasks from being submitted to the pool"""
#         self.add = RaiseException("This %s has been closed" % self.__class__.__name__)

    def join(self, timeout=None, raise_error=False):
        timeout = Timeout.start_new(timeout)
        try:
            try:
                while self.greenlets:
                    joinall(self.greenlets, raise_error=raise_error)
            except Timeout, ex:
                if ex is not timeout:
                    raise
        finally:
            timeout.cancel()

    def kill(self, exception=GreenletExit, block=False, timeout=None):
        timer = Timeout.start_new(timeout)
        try:
            while self.greenlets:
                for greenlet in self.greenlets:
                    if greenlet not in self.dying:
                        greenlet.kill(exception)
                        self.dying.add(greenlet)
                if not block:
                    break
                joinall(self.greenlets)
        finally:
            timer.cancel()

    def killone(self, greenlet, exception=GreenletExit, block=False, timeout=None):
        if greenlet not in self.dying and greenlet in self.greenlets:
            greenlet.kill(exception)
            self.dying.add(greenlet)
            if block:
                greenlet.join(timeout)

    def apply(self, func, args=None, kwds=None):
        """Equivalent of the apply() builtin function. It blocks till the result is ready."""
        if args is None:
            args = ()
        if kwds is None:
            kwds = {}
        if getcurrent() in self:
            return func(*args, **kwds)
        else:
            return self.spawn(func, *args, **kwds).get()

    def apply_async(self, func, args=None, kwds=None, callback=None):
        """A variant of the apply() method which returns a Greenlet object.

        If callback is specified then it should be a callable which accepts a single argument. When the result becomes ready
        callback is applied to it (unless the call failed)."""
        if args is None:
            args = ()
        if kwds is None:
            kwds = {}
        greenlet = self.spawn(func, *args, **kwds)
        if callback is not None:
            greenlet.link(pass_value(callback))
        return greenlet

    def map(self, func, iterable):
        greenlets = [self.spawn(func, item) for item in iterable]
        return [greenlet.get() for greenlet in greenlets]

    def map_async(self, func, iterable, callback=None):
        """
        A variant of the map() method which returns a Greenlet object.

        If callback is specified then it should be a callable which accepts a
        single argument.
        """
        greenlets = [self.spawn(func, item) for item in iterable]
        result = self.spawn(get_values, greenlets)
        if callback is not None:
            result.link(pass_value(callback))
        return result

    def imap(self, func, iterable):
        """An equivalent of itertools.imap()"""
        greenlets = [self.spawn(func, item) for item in iterable]
        for greenlet in greenlets:
            yield greenlet.get()

    def imap_unordered(self, func, iterable):
        """The same as imap() except that the ordering of the results from the
        returned iterator should be considered arbitrary."""
        from gevent.queue import Queue
        q = Queue()
        greenlets = [self.spawn(func, item) for item in iterable]
        for greenlet in greenlets:
            greenlet.rawlink(q.put)
        for _ in xrange(len(greenlets)):
            yield q.get().get()

    def full(self):
        return False


class Pool(GreenletSet):

    def __init__(self, size=None):
        if size is not None and size < 0:
            raise ValueError('Invalid size for pool (positive integer or None required): %r' % (size, ))
        GreenletSet.__init__(self)
        self.size = size
        self.waiting = deque()

    def full(self):
        return self.free_count() <= 0

    def free_count(self):
        if self.size is None:
            return 1
        return max(0, self.size - len(self) - len(self.waiting))

    def start(self, greenlet):
        if self.size is not None and len(self) >= self.size:
            self.waiting.append(greenlet)
        else:
            greenlet.start()
            self.add(greenlet)

    def spawn(self, function, *args, **kwargs):
        greenlet = Greenlet(function, *args, **kwargs)
        self.start(greenlet)
        return greenlet

    def discard(self, greenlet):
        GreenletSet.discard(self, greenlet)
        while self.waiting and len(self) < self.size:
            greenlet = self.waiting.popleft()
            greenlet.start()
            self.add(greenlet)

    def kill(self, exception=GreenletExit, block=False, timeout=None):
        for greenlet in self.waiting:
            greenlet.kill(exception)
        self.waiting.clear()
        return GreenletSet.kill(self, exception=exception, block=block, timeout=timeout)


def get_values(greenlets):
    joinall(greenlets)
    return [x.value for x in greenlets]


class pass_value(object):
    __slots__ = ['callback']

    def __init__(self, callback):
        self.callback = callback

    def __call__(self, source):
        if source.successful():
            self.callback(source.value)

    def __hash__(self):
        return hash(self.callback)

    def __eq__(self, other):
        return self.callback == getattr(other, 'callback', other)

    def __str__(self):
        return str(self.callback)

    def __repr__(self):
        return repr(self.callback)

    def __getattr__(self, item):
        assert item != 'callback'
        return getattr(self.callback, item)