r""" Utilities for rxcpp  Observables in python"""

import asyncio


class QueueIterator:
    """An Iterator for an asyncio.Queue object

    Allows a consumer to iterate the elements of a queue using an `async for ` clause.
    Producer may push elements through the queue object supplied to the constructor.
    `done` should be set once the producer stops pushing elements. In that case the
    iterator will consume any elements in the queue and exit.

    The producer may use `queue.join()` to monitor when the queue is empty
    """

    def __init__(self, queue, done):
        if not isinstance(queue, asyncio.Queue):
            raise TypeError("first argument is not of type asyncio.Queue")

        if not isinstance(done, asyncio.Future):
            raise TypeError("second argument is not of type asyncio.Future")

        self._Queue = queue
        self._future = done

    def get_queue(self):
        return self._Queue

    def get_future(self):
        return self._future

    def __aiter__(self):
        return self

    async def __anext__(self):
        while True:
            if self._future.cancelled():
                raise StopAsyncIteration
            if self._future.done() and self._future.exception() is not None:
                raise self._future.exception()
            if self._Queue.qsize() == 0 and self._future.done():
                raise StopAsyncIteration
            if self._Queue.qsize() > 0:
                f = self._Queue.get_nowait()
                self._Queue.task_done()
                return f
            else:
                # Wait until an element is pushed or future is set
                task = asyncio.create_task(self._Queue.get())
                done, _ = await asyncio.wait(
                    {task, self._future},
                    return_when=asyncio.FIRST_COMPLETED,
                )
                if task in done:
                    self._Queue.task_done()
                    return task.result()
                else:
                    task.cancel()
