"""ASR engine implementation using faster-whisper.

Provides Whisper-based transcription with word-level timestamps.
"""

from __future__ import annotations

import asyncio
from collections.abc import Iterable, Iterator
from typing import TYPE_CHECKING, Final, Protocol, cast

from noteflow.infrastructure.logging import get_logger, log_timing

if TYPE_CHECKING:
    import numpy as np
    from numpy.typing import NDArray


class _WhisperWord(Protocol):
    word: str
    start: float
    end: float
    probability: float


class _WhisperSegment(Protocol):
    text: str
    start: float
    end: float
    words: Iterable[_WhisperWord] | None
    avg_logprob: float
    no_speech_prob: float


class _WhisperInfo(Protocol):
    language: str
    language_probability: float


class _WhisperModel(Protocol):
    def transcribe(
        self,
        audio: NDArray[np.float32],
        *,
        language: str | None = None,
        word_timestamps: bool = ...,
        beam_size: int = ...,
        vad_filter: bool = ...,
    ) -> tuple[Iterable[_WhisperSegment], _WhisperInfo]: ...

from noteflow.infrastructure.asr.dto import AsrResult, WordTiming

logger = get_logger(__name__)

# Available model sizes
VALID_MODEL_SIZES: Final[tuple[str, ...]] = (
    "tiny",
    "tiny.en",
    "base",
    "base.en",
    "small",
    "small.en",
    "medium",
    "medium.en",
    "large-v1",
    "large-v2",
    "large-v3",
)


class FasterWhisperEngine:
    """faster-whisper based ASR engine.

    Uses CTranslate2 for efficient Whisper inference on CPU or GPU.
    """

    def __init__(
        self,
        compute_type: str = "int8",
        device: str = "cpu",
        num_workers: int = 1,
    ) -> None:
        """Initialize the engine.

        Args:
            compute_type: Computation type ("int8", "float16", "float32").
            device: Device to use ("cpu" or "cuda").
            num_workers: Number of worker threads.
        """
        self._compute_type = compute_type
        self._device = device
        self._num_workers = num_workers
        self._model: _WhisperModel | None = None
        self._model_size: str | None = None

    def load_model(self, model_size: str = "base") -> None:
        """Load the ASR model.

        Args:
            model_size: Model size (e.g., "tiny", "base", "small").

        Raises:
            ValueError: If model_size is invalid.
            RuntimeError: If model loading fails.
        """
        from faster_whisper import WhisperModel

        if model_size not in VALID_MODEL_SIZES:
            raise ValueError(
                f"Invalid model size: {model_size}. Valid sizes: {', '.join(VALID_MODEL_SIZES)}"
            )

        with log_timing(
                "asr_model_load",
                model_size=model_size,
                device=self._device,
                compute_type=self._compute_type,
            ):
            try:
                model = WhisperModel(
                    model_size,
                    device=self._device,
                    compute_type=self._compute_type,
                    num_workers=self._num_workers,
                )
                self._model = cast(_WhisperModel, model)
                self._model_size = model_size
            except (RuntimeError, OSError, ValueError) as e:
                raise RuntimeError(f"Failed to load model: {e}") from e

        logger.info(
            "asr_model_loaded",
            model_size=model_size,
            device=self._device,
            compute_type=self._compute_type,
        )

    def transcribe(
        self,
        audio: NDArray[np.float32],
        language: str | None = None,
    ) -> Iterator[AsrResult]:
        """Transcribe audio and yield results.

        Args:
            audio: Audio samples as float32 array (16kHz mono, normalized).
            language: Optional language code (e.g., "en").

        Yields:
            AsrResult segments with word-level timestamps.
        """
        if self._model is None:
            raise RuntimeError("Model not loaded. Call load_model() first.")

        # Transcribe with word timestamps
        segments, info = self._model.transcribe(
            audio,
            language=language,
            word_timestamps=True,
            beam_size=5,
            vad_filter=True,  # Filter out non-speech
        )

        logger.debug(
            "Detected language: %s (prob: %.2f)",
            info.language,
            info.language_probability,
        )

        for segment in segments:
            # Convert word info to WordTiming objects
            words: list[WordTiming] = []
            if segment.words:
                words = [
                    WordTiming(
                        word=word.word,
                        start=word.start,
                        end=word.end,
                        probability=word.probability,
                    )
                    for word in segment.words
                ]

            yield AsrResult(
                text=segment.text.strip(),
                start=segment.start,
                end=segment.end,
                words=tuple(words),
                language=info.language,
                language_probability=info.language_probability,
                avg_logprob=segment.avg_logprob,
                no_speech_prob=segment.no_speech_prob,
            )

    async def transcribe_async(
        self,
        audio: NDArray[np.float32],
        language: str | None = None,
    ) -> list[AsrResult]:
        """Transcribe audio asynchronously using executor.

        Offloads blocking transcription to a thread pool executor to avoid
        blocking the asyncio event loop.

        Args:
            audio: Audio samples as float32 array (16kHz mono, normalized).
            language: Optional language code (e.g., "en").

        Returns:
            List of AsrResult segments with word-level timestamps.
        """
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(
            None,
            self._transcribe_to_list,
            audio,
            language,
        )

    def _transcribe_to_list(
        self,
        audio: NDArray[np.float32],
        language: str | None,
    ) -> list[AsrResult]:
        return list(self.transcribe(audio, language))

    @property
    def is_loaded(self) -> bool:
        """Return True if model is loaded."""
        return self._model is not None

    @property
    def model(self) -> _WhisperModel | None:
        """Return the loaded model instance, if available."""
        return self._model

    @property
    def model_size(self) -> str | None:
        """Return the loaded model size, or None if not loaded."""
        return self._model_size

    def unload(self) -> None:
        """Unload the model to free memory."""
        self._model = None
        self._model_size = None
        logger.info("Model unloaded")

    @property
    def compute_type(self) -> str:
        """Return the compute type."""
        return self._compute_type

    @property
    def device(self) -> str:
        """Return the device."""
        return self._device
