import abc
import sys

if sys.platform == "darwin":
    import mlx.core as mx
    import mlx_whisper
    from mlx_whisper.transcribe import ModelHolder

import numpy as np
import torch
import whisper


class SpeechRecognizer(abc.ABC):
    def __init__(self, limit_output_length=True):
        """
        :param limit_output_length: When `True`, the length of the generated speech will be limited
        by the length of the input audio and some heuristics.
        """
        self.limit_output_length = limit_output_length

    @abc.abstractmethod
    def load_model(self): ...

    @abc.abstractmethod
    def recognize_speech(self, audio: np.ndarray) -> str:
        """
        Recognize speech from the given audio sample.

        :param audio: A full utterance sample. Audio must be 16 kHz, mono, np.float32, values in the
        range [-1.0, 1.0].
        :return: Recognized speech.
        """

    @staticmethod
    def _estimate_max_tokens(audio: np.ndarray) -> int:
        """
        Estimate the maximum length of a given audio sample in tokens. Assumes a maximum speaking
        rate of 300 words per minute (2x average), and assumes that 3 words is 4 tokens.

        :param audio: The audio sample (16 kHz) to use for length estimation.
        :return: The estimated length of the transcribed audio in tokens.
        """
        length_seconds = len(audio) / 16_000
        length_minutes = length_seconds / 60
        word_count = length_minutes * 300
        token_count = word_count / 3 * 4
        return int(token_count)

    def _get_decode_options(self, audio: np.ndarray) -> dict:
        """
        :param audio: The audio sample (16 kHz) to use to determine options like max decode length.
        :return: A dict that can be used to construct `whisper.DecodingOptions`.
        """
        options = {}
        if self.limit_output_length:
            options["sample_len"] = self._estimate_max_tokens(audio)
        return options

    @staticmethod
    def best_type():
        """Get the best type of SpeechRecognizer based on system capabilities."""
        if torch.mps.is_available():
            print("Choosing MLX Whisper model.")
            return MLXWhisperSpeechRecognizer()
        else:
            print("Choosing reference Whisper model.")
            return OpenAIWhisperSpeechRecognizer()


class MLXWhisperSpeechRecognizer(SpeechRecognizer):
    def __init__(self, limit_output_length=True):
        super().__init__(limit_output_length)
        self.was_loaded = False
        self.model_name = "mlx-community/whisper-small.en-mlx"

    def load_model(self):
        if self.was_loaded:
            return
        # There appears to be no dedicated mechanism to preload a model, but this `get_model` does
        # store it in memory for later usage
        ModelHolder.get_model(self.model_name, mx.float16)
        self.was_loaded = True

    def recognize_speech(self, audio: np.ndarray) -> str:
        self.load_model()
        return mlx_whisper.transcribe(
            audio, path_or_hf_repo=self.model_name, decode_options=self._get_decode_options(audio)
        )["text"]
        return mlx_whisper.transcribe(audio, path_or_hf_repo=self.model_name)["text"].strip()


class OpenAIWhisperSpeechRecognizer(SpeechRecognizer):
    def __init__(self, limit_output_length=True):
        super().__init__(limit_output_length)
        self.model = None

    def load_model(self):
        if self.model is not None:
            return
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model = whisper.load_model("small.en", device=device)

    def recognize_speech(self, audio: np.ndarray) -> str:
        self.load_model()
        return whisper.transcribe(
            self.model, audio, decode_options=self._get_decode_options(audio)
        )["text"]
