from __future__ import annotations

import asyncio
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from typing import Protocol

import numpy as np


@dataclass(frozen=True)
class Transcript:
    text: str
    language: str
    probability: float


class Transcriber(Protocol):
    async def warmup(self) -> None:
        ...

    async def transcribe(
        self,
        samples: np.ndarray,
        *,
        final: bool,
    ) -> Transcript:
        ...


class FasterWhisperTranscriber:
    def __init__(self, model_dir: Path, compute_type: str) -> None:
        self._model_dir = model_dir
        self._compute_type = compute_type
        self._executor = ThreadPoolExecutor(
            max_workers=1,
            thread_name_prefix="dictation",
        )
        self._model = None
        self._load_lock = asyncio.Lock()

    async def warmup(self) -> None:
        if self._model is not None:
            return
        async with self._load_lock:
            if self._model is not None:
                return
            loop = asyncio.get_running_loop()
            self._model = await loop.run_in_executor(
                self._executor,
                self._load_and_warmup,
            )

    def _load_and_warmup(self):
        from faster_whisper import WhisperModel

        if not (self._model_dir / "model.bin").is_file():
            raise FileNotFoundError(
                f"Model not found at {self._model_dir}. "
                "Run: bazel run //dictation:download_model"
            )
        model = WhisperModel(
            str(self._model_dir),
            device="cuda",
            compute_type=self._compute_type,
            local_files_only=True,
        )
        # Loading weights does not initialize all CUDA kernels. Execute and
        # consume one short silent inference now so the user's first utterance
        # does not pay the one-time GPU setup cost.
        segments, _ = model.transcribe(
            np.zeros(8000, dtype=np.float32),
            beam_size=1,
            best_of=1,
            condition_on_previous_text=False,
            vad_filter=False,
            without_timestamps=True,
        )
        list(segments)
        return model

    async def transcribe(
        self,
        samples: np.ndarray,
        *,
        final: bool,
    ) -> Transcript:
        await self.warmup()
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(
            self._executor,
            self._transcribe_sync,
            np.asarray(samples, dtype=np.float32),
            final,
        )

    def _transcribe_sync(
        self,
        samples: np.ndarray,
        final: bool,
    ) -> Transcript:
        segments, info = self._model.transcribe(
            samples,
            beam_size=1,
            best_of=1,
            condition_on_previous_text=False,
            vad_filter=False,
            without_timestamps=True,
        )
        text = "".join(segment.text for segment in segments).strip()
        return Transcript(
            text=text,
            language=info.language or "",
            probability=float(info.language_probability or 0.0),
        )

    def close(self) -> None:
        self._executor.shutdown(wait=False, cancel_futures=True)