from __future__ import annotations

import asyncio
from concurrent.futures import ThreadPoolExecutor
from typing import Any, AsyncIterator

from audio.pipeline import AudioPipeline

_SENTINEL = object()


class GpuSerialQueue:
    """Bounded replica pool. Each replica runs batched Whisper on one file."""

    def __init__(self, pipeline: AudioPipeline, max_queue: int, gpu_concurrency: int) -> None:
        self.pipeline = pipeline
        workers = max(1, gpu_concurrency)
        self._gpu = asyncio.Semaphore(workers)
        self._slots = asyncio.Semaphore(max_queue)
        self._pool = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="asr")

    async def run(self, **kwargs: Any) -> dict:
        try:
            await asyncio.wait_for(self._slots.acquire(), timeout=0.05)
        except TimeoutError as exc:
            raise QueueFullError("audio GPU queue is full") from exc
        try:
            async with self._gpu:
                loop = asyncio.get_running_loop()
                return await loop.run_in_executor(
                    self._pool, lambda: self.pipeline.run(**kwargs)
                )
        finally:
            self._slots.release()

    async def stream(self, **kwargs: Any) -> AsyncIterator[dict[str, Any]]:
        try:
            await asyncio.wait_for(self._slots.acquire(), timeout=0.05)
        except TimeoutError as exc:
            raise QueueFullError("audio GPU queue is full") from exc
        await self._gpu.acquire()
        loop = asyncio.get_running_loop()
        out: asyncio.Queue[Any] = asyncio.Queue()

        def produce() -> None:
            try:
                for event in self.pipeline.run_events(**kwargs):
                    asyncio.run_coroutine_threadsafe(out.put(event), loop).result()
                asyncio.run_coroutine_threadsafe(out.put(_SENTINEL), loop).result()
            except Exception as exc:
                asyncio.run_coroutine_threadsafe(out.put(exc), loop).result()

        try:
            self._pool.submit(produce)
            while True:
                item = await out.get()
                if item is _SENTINEL:
                    break
                if isinstance(item, Exception):
                    raise item
                yield item
        finally:
            self._gpu.release()
            self._slots.release()


class QueueFullError(RuntimeError):
    pass