"""Whisper STT implementation using HuggingFace transformers."""

from __future__ import annotations

import numpy as np
from roohai.base import STTModel


class WhisperSTT(STTModel):
    META = {
        "display_name": "OpenAI Whisper",
        "description": "Open-source speech recognition model by OpenAI. Accurate, multilingual, runs locally.",
        "kind": "local",
        "pip_dependencies": ["transformers>=4.46.0,<5.0.0", "torch", "sentencepiece"],
        "hf_models": ["openai/whisper-tiny"],
        "estimated_download_mb": 150,
        "config_fields": [
            {
                "name": "model_id",
                "label": "Model ID",
                "type": "text",
                "required": False,
                "default": "openai/whisper-tiny",
            },
            {
                "name": "language",
                "label": "Language",
                "type": "text",
                "required": False,
                "default": "en",
            },
        ],
    }

    def __init__(self, model_id: str = "openai/whisper-tiny", language: str = "en") -> None:
        self.model_id = model_id
        self.language = language
        self._processor = None
        self._model = None

    def load(self) -> None:
        from transformers import WhisperProcessor, WhisperForConditionalGeneration

        self._processor = WhisperProcessor.from_pretrained(self.model_id)
        self._model = WhisperForConditionalGeneration.from_pretrained(self.model_id)

    def transcribe(self, audio: np.ndarray, sample_rate: int = 16000) -> str:
        if not self.is_loaded:
            raise RuntimeError("Model not loaded. Call load() first.")
        input_features = self._processor(
            audio,
            sampling_rate=sample_rate,
            return_tensors="pt",
            language=self.language,
        ).input_features
        predicted_ids = self._model.generate(input_features)
        transcription = self._processor.batch_decode(predicted_ids, skip_special_tokens=True)
        return transcription[0].strip()

    def unload(self) -> None:
        if self._model is not None:
            from roohai.base import _release_torch_memory
            del self._model
            del self._processor
            self._model = None
            self._processor = None
            _release_torch_memory()

    @property
    def is_loaded(self) -> bool:
        return self._model is not None
