"""Wav2Vec2 STT implementation using HuggingFace transformers."""

from __future__ import annotations

import numpy as np
import torch
from roohai.base import STTModel


class Wav2Vec2STT(STTModel):
    META = {
        "display_name": "Wav2Vec2",
        "description": "Facebook's self-supervised speech recognition model. English-only, runs locally.",
        "kind": "local",
        "pip_dependencies": ["transformers>=4.46.0,<5.0.0", "torch"],
        "hf_models": ["facebook/wav2vec2-base-960h"],
        "estimated_download_mb": 360,
        "config_fields": [
            {
                "name": "model_id",
                "label": "Model ID",
                "type": "text",
                "required": False,
                "default": "facebook/wav2vec2-base-960h",
            },
        ],
    }

    def __init__(self, model_id: str = "facebook/wav2vec2-base-960h") -> None:
        self.model_id = model_id
        self._processor = None
        self._model = None

    def load(self) -> None:
        from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC

        self._processor = Wav2Vec2Processor.from_pretrained(self.model_id)
        self._model = Wav2Vec2ForCTC.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.")
        inputs = self._processor(
            audio, sampling_rate=sample_rate, return_tensors="pt", padding=True
        )
        with torch.no_grad():
            logits = self._model(**inputs).logits
        predicted_ids = torch.argmax(logits, dim=-1)
        transcription = self._processor.batch_decode(predicted_ids)
        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
