import logging
from typing import Iterator, Optional

from huggingface_hub import constants as hf_constants
from transformers import AutomaticSpeechRecognitionPipeline, pipeline, \
    ParakeetFeatureExtractor, ParakeetForTDT, ParakeetTokenizer

from sttts.api.message import Message, ModelError, ModelNotFoundError, SttKeywordMap, SttMessageType
from sttts.api.model import Recognizer, RecognizerModel
from sttts.utils.utils import MathUtils, PerfCounter


class ParakeetRecognizer(RecognizerModel, Recognizer):
    """
    Use the `Nvidia Parakeet <https://catalog.ngc.nvidia.com/orgs/nvidia/-/collections/parakeet-tdt-0.6b/->`__
    `model <https://huggingface.co/nvidia/parakeet-tdt-0.6b-v3>`__ with transformers
    `pipeline <https://huggingface.co/docs/transformers/en/model_doc/parakeet>`__ defaults.
    Note that (in contrast to `Canary <https://huggingface.co/nvidia/canary-1b-v2>`__), the spoken language is
    guessed and cannot be fixed in advance.
    """

    def __init__(self, keywords: SttKeywordMap, *,
                 model_name: str = "nvidia/parakeet-tdt-0.6b-v3", download: bool = False, device: Optional[str] = None,
                 **kwargs) -> None:
        """
        :param str model_name: HuggingFace model name to use, ``nvidia/parakeet-tdt-0.6b-v3`` per default.
        :param bool download: Opt-in automatic download. Otherwise, the models must exist in offline mode.
        :param str device: Device to explicitly use, for example ``cuda`` if available or ``cpu``.
        """

        self._logger: logging.Logger = logging.getLogger(self.__class__.__name__)
        self._sample_rate: int = 16000
        self._keywords: SttKeywordMap = keywords

        if not download:
            # XXX: pipeline accepts local_files_only, but mixes hub and model kwargs, leading to
            #      ValueError: The following `model_kwargs` are not used by the model: ['local_files_only']
            hf_constants.HF_HUB_OFFLINE = True
        try:
            self._model_name: str = model_name
            self._model: AutomaticSpeechRecognitionPipeline = pipeline(
                "automatic-speech-recognition", model=self._model_name, device=device, **kwargs,
            )
        except OSError as e:
            raise ModelNotFoundError(self.__class__.__name__, model_name, str(e), None) from None
        except Exception as e:
            raise ModelError(self.__class__.__name__, model_name, repr(e))
        else:
            assert isinstance(self._model, AutomaticSpeechRecognitionPipeline)
            assert isinstance(self._model.model, ParakeetForTDT)
            assert isinstance(self._model.feature_extractor, ParakeetFeatureExtractor)
            assert isinstance(self._model.tokenizer, ParakeetTokenizer)

    def sample_rate(self) -> int:
        return self._sample_rate

    def accept(self, buffer: bytes) -> Iterator[Message[SttMessageType]]:
        with PerfCounter(self._logger, logging.INFO, "msec") as counter:
            utterance: str = self._model(MathUtils.buf2arr(buffer))["text"]  # type: ignore
            counter(round(len(buffer) / 2.0 / self._sample_rate * 1000.0))
            yield from self._keywords.map_utterance(utterance)

    def __enter__(self) -> Recognizer:
        self._logger.info(f"Entering <{self._model.__class__.__name__}>[{self._model_name}]")
        self._logger.debug(str(self._model))
        return self

    def __exit__(self, *args) -> None:
        pass