import asyncio
from typing import AsyncIterable, Coroutine

from aiospb.messages.values import Message, MQTTMessage


class SubscriptionObserver:
    def __init__(self, subscription: AsyncIterable[Message]):
        self._queue = asyncio.Queue()

        async def put_messages():
            async for msg in subscription:
                self._queue.put_nowait(msg)

        self._task = asyncio.create_task(put_messages())

    async def __aenter__(self):
        return self._queue

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        self._task.cancel()
        try:
            await self._task
        except asyncio.CancelledError:
            ...


class Stack:
    def __init__(self):
        self._queue = asyncio.Queue()

    def put(self, msg):
        self._queue.put_nowait(msg)

    def __aiter__(self):
        return self

    async def __anext__(self):
        msg = await self._queue.get()
        if isinstance(msg, Exception):
            raise msg
        return msg
