from __future__ import annotations

from fastapi.responses import StreamingResponse

from app.modules.api.application.stream_service import StreamService


class StreamController:
    def __init__(self, service: StreamService) -> None:
        self._service = service

    async def stream(self, request_id: str) -> StreamingResponse:
        queue = await self._service.subscribe(request_id)

        async def event_stream():
            import asyncio

            heartbeat = 10
            try:
                while True:
                    try:
                        event = await asyncio.wait_for(queue.get(), timeout=heartbeat)
                        yield self._service.encode(event)
                    except asyncio.TimeoutError:
                        yield ": keepalive\n\n"
            finally:
                await self._service.unsubscribe(request_id, queue)

        return StreamingResponse(
            event_stream(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache, no-transform",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",
            },
        )
