"""NVIDIA NeMo STT implementations (Canary-Qwen, Parakeet)."""

from __future__ import annotations

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

logger = logging.getLogger(__name__)


class NvidiaCanarySTT(STTModel):
    """NVIDIA Canary-Qwen 2.5B speech model using NeMo SALM."""

    META = {
        "display_name": "NVIDIA Canary-Qwen 2.5B",
        "description": "NVIDIA Canary-Qwen 2.5B speech model. Multilingual, requires CUDA GPU.",
        "kind": "local",
        "pip_dependencies": ["nemo_toolkit[asr]", "torch", "soundfile"],
        "estimated_download_mb": 5000,
        "config_fields": [
            {
                "name": "model_id",
                "label": "Model ID",
                "type": "text",
                "required": False,
                "default": "nvidia/canary-qwen-2.5b",
            },
        ],
    }

    def __init__(self, model_id: str = "nvidia/canary-qwen-2.5b") -> None:
        self.model_id = model_id
        self._model = None

    def load(self) -> None:
        import torch
        from nemo.collections.speechlm2.models import SALM

        logger.info("Loading NVIDIA Canary-Qwen model: %s", self.model_id)
        self._model = SALM.from_pretrained(self.model_id)
        if torch.cuda.is_available():
            self._model = self._model.cuda()
            logger.info("Canary-Qwen model moved to CUDA")
        else:
            logger.warning("CUDA not available — Canary-Qwen will run on CPU (very slow)")

    def transcribe(self, audio: np.ndarray, sample_rate: int = 16000) -> str:
        if not self.is_loaded:
            raise RuntimeError("Model not loaded. Call load() first.")

        import tempfile
        import soundfile as sf

        # SALM expects an audio file path, so write a temp WAV file
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            sf.write(f.name, audio, sample_rate)
            tmp_path = f.name

        try:
            result = self._model.transcribe([tmp_path])
            # SALM.transcribe returns a list of transcriptions
            if isinstance(result, list) and len(result) > 0:
                text = result[0]
                if isinstance(text, dict):
                    text = text.get("text", str(text))
                return str(text).strip()
            return ""
        finally:
            import os
            os.unlink(tmp_path)

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

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


class NvidiaParakeetSTT(STTModel):
    """NVIDIA Parakeet TDT ASR model using NeMo."""

    META = {
        "display_name": "NVIDIA Parakeet TDT 0.6B V2",
        "description": "NVIDIA Parakeet TDT 0.6B V2 ASR model. Fast, accurate, supports timestamps. Requires CUDA GPU.",
        "kind": "local",
        "pip_dependencies": ["nemo_toolkit[asr]", "torch", "soundfile"],
        "estimated_download_mb": 2400,
        "config_fields": [
            {
                "name": "model_id",
                "label": "Model ID",
                "type": "text",
                "required": False,
                "default": "nvidia/parakeet-tdt-0.6b-v2",
            },
        ],
    }

    def __init__(self, model_id: str = "nvidia/parakeet-tdt-0.6b-v2") -> None:
        self.model_id = model_id
        self._model = None

    def load(self) -> None:
        import torch
        import nemo.collections.asr as nemo_asr

        logger.info("Loading NVIDIA Parakeet model: %s", self.model_id)
        self._model = nemo_asr.models.ASRModel.from_pretrained(model_name=self.model_id)
        if torch.cuda.is_available():
            self._model = self._model.cuda()
            logger.info("Parakeet model moved to CUDA")
        else:
            logger.warning("CUDA not available — Parakeet will run on CPU (very slow)")

    def transcribe(self, audio: np.ndarray, sample_rate: int = 16000) -> str:
        if not self.is_loaded:
            raise RuntimeError("Model not loaded. Call load() first.")

        import tempfile
        import soundfile as sf

        # NeMo ASR models expect file paths
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            sf.write(f.name, audio, sample_rate)
            tmp_path = f.name

        try:
            result = self._model.transcribe([tmp_path])
            if isinstance(result, list) and len(result) > 0:
                text = result[0]
                if isinstance(text, dict):
                    text = text.get("text", str(text))
                return str(text).strip()
            return ""
        finally:
            import os
            os.unlink(tmp_path)

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

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