import logging
from typing import Iterator, Optional, Tuple

from datasets import Dataset, load_dataset, config as datasets_config
from huggingface_hub import constants as hf_constants
from torch import tensor, FloatTensor, LongTensor, Tensor
from transformers import SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5Processor

from sttts.api.message import ModelError, ModelNotFoundError
from sttts.api.model import Synthesizer, SynthesizerModel
from sttts.utils.utils import MathUtils, PerfCounter


class _TransformersSynthesizer(Synthesizer):
    def __init__(self, processor: SpeechT5Processor, model: SpeechT5ForTextToSpeech,
                 vocoder: SpeechT5HifiGan, speaker: Tensor, device: Optional[str]) -> None:
        self._logger: logging.Logger = logging.getLogger(self.__class__.__name__.lstrip("_"))
        self._processor: SpeechT5Processor = processor
        self._model: SpeechT5ForTextToSpeech = model.to(device)  # type: ignore
        self._vocoder: SpeechT5HifiGan = vocoder.to(device)  # type: ignore
        self._speaker: FloatTensor = speaker.to(device)  # type: ignore
        self._device: Optional[str] = device

    def sample_rate(self) -> int:
        return 16000

    def generate(self, utterance: str) -> Iterator[bytes]:
        with PerfCounter(self._logger, logging.INFO, "msec") as counter:
            inputs: LongTensor = self._processor(text=utterance, return_tensors="pt")["input_ids"].to(self._device)
            speech: FloatTensor = self._model.generate_speech(  # type: ignore
                inputs, self._speaker, vocoder=self._vocoder,
            )
            counter(round(len(speech) / self.sample_rate() * 1000.0))
            buffer: bytes = MathUtils.arr2buf(speech.cpu().numpy())
        yield buffer


class TransformersSynthesizerModel(SynthesizerModel):
    """
    Use the `SpeechT5 <https://huggingface.co/blog/speecht5>`__ Hugging Face TTS
    `transformers <https://huggingface.co/microsoft/speecht5_tts>`__ model.
    The vocoder comes from a `fork <https://huggingface.co/datasets/regisss/cmu-arctic-xvectors>`__ of the
    `CMU ARCTIC <https://huggingface.co/datasets/Matthijs/cmu-arctic-xvectors>`__ speaker embeddings.
    """

    def __init__(self, *, model_name: str, download: bool = False, device: Optional[str] = None) -> None:
        """
        :param str model_name: Model from ``cmu-arctic-xvectors``, for example ``cmu_us_slt_arctic-wav-arctic_a0001``.
        :param bool download: Opt-in automatic download. Otherwise, the models and datasets 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._pipeline: Tuple[str, str, str] = ("microsoft/speecht5_tts", "microsoft/speecht5_tts",
                                                "microsoft/speecht5_hifigan")
        self._model_dataset: str = "regisss/cmu-arctic-xvectors"  # "Matthijs/cmu-arctic-xvectors"
        self._model_name: str = model_name
        self._offline: bool = not download
        self._device: Optional[str] = device

        if self._offline:  # no local_files_only flag for datasets; from_pretrained still checks metadata
            datasets_config.HF_HUB_OFFLINE = True
            datasets_config.HF_DATASETS_OFFLINE = True
            hf_constants.HF_HUB_OFFLINE = True

        self._speaker_model: Tensor = self._get_speaker_model(self._model_dataset, self._model_name)

    @classmethod
    def _get_speaker_model(cls, dataset: str, name: str) -> Tensor:
        try:
            embeddings: Dataset = load_dataset(dataset, split="validation")
            model_idx: int = embeddings["filename"].index(name)
            return tensor(embeddings[model_idx]["xvector"]).unsqueeze(0)
        except (OSError, KeyError, ValueError) as e:
            raise ModelNotFoundError(cls.__name__, f"{dataset}/{name}", repr(e), None) from None

    def __enter__(self) -> Synthesizer:
        try:
            return _TransformersSynthesizer(
                SpeechT5Processor.from_pretrained(self._pipeline[0], local_files_only=self._offline,
                                                  device=self._device),
                SpeechT5ForTextToSpeech.from_pretrained(self._pipeline[1], local_files_only=self._offline),
                SpeechT5HifiGan.from_pretrained(self._pipeline[2], local_files_only=self._offline),
                self._speaker_model,
                self._device,
            )
        except OSError as e:
            raise ModelNotFoundError(self.__class__.__name__, self._pipeline[0], str(e), None) from None
        except Exception as e:
            raise ModelError(self.__class__.__name__, self._pipeline[0], repr(e))

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