"""
Wrapper module for asynchronous task handling.
"""
import asyncio, inspect


class UpdateTask(asyncio.Task):
    """Wrap a coroutine in a future"""

    def __init__(self, coro, desc, loop=None):
        """
        Arguments:
            - coro: awaitable object
            - desc<tuple>: (Resource, key)
        """
        assert asyncio.iscoroutine(coro), coro
        self._desc = desc
        super().__init__(coro)

    def __repr__(self):
        res, pk = self._desc
        return "<UpdateTask: ({}, {})>".format(res.tag, pk)


def gather(jobs):
    "Aggregate and collect jobs"
    return asyncio.gather(*jobs, return_exceptions=False)


def wrap_generator(func):
    """
    Decorator to convert a generator function to an async function which collects
    and returns generator results, returning a list if there are multiple results
    """

    async def _wrapped(*a, **k):
        r, ret = None, []
        gen = func(*a, **k)
        while True:
            try:
                item = gen.send(r)
            except StopIteration:
                break
            if inspect.isawaitable(item):
                r = await item
            else:
                r = item
            ret.append(r)

        if len(ret) == 1:
            return ret.pop()
        return ret

    return _wrapped


def run_task(func):
    """
    Decorator to wrap an async function in an event loop.
    Use for main sync interface methods.
    """

    def _wrapped(*a, **k):
        loop = asyncio.get_event_loop()
        return loop.run_until_complete(func(*a, **k))

    return _wrapped
