import asyncio
from collections.abc import AsyncGenerator

from judy_core.events import OutputEvent, TaskFailResult, TaskSuccessResult
from judy_core.topics import OUTPUT_TOPIC_KIND
from tongcell.core import (
    ActorStopMessage,
    RoutedActor,
    message_handler,
    subscribe_topic,
)


@subscribe_topic(OUTPUT_TOPIC_KIND)
class OutputStreamer(RoutedActor):
    def __init__(self, description: str = "", **kwargs) -> None:
        super().__init__(description, **kwargs)
        self._events = asyncio.Queue[OutputEvent]()

    @message_handler()
    async def on_stop(self, m: ActorStopMessage):
        # self._events.shutdown()
        pass

    @message_handler()
    async def _receive_output_event(self, event: OutputEvent):
        await self._events.put(event)

    async def stream(self) -> AsyncGenerator[OutputEvent]:
        try:
            while True:
                event = await self._events.get()
                yield event
                if isinstance(event, TaskSuccessResult | TaskFailResult):
                    return
        except asyncio.QueueShutDown:
            return
