import warnings
from typing import Optional, Union

import numpy as np
import torch

from ..audio import convert_demucs_kwargs, prep_audio
from ..non_whisper import transcribe_any
from ..utils import isolate_useful_options


HF_MODELS = {
    "tiny.en": "openai/whisper-tiny.en",
    "tiny": "openai/whisper-tiny",
    "base.en": "openai/whisper-base.en",
    "base": "openai/whisper-base",
    "small.en": "openai/whisper-small.en",
    "small": "openai/whisper-small",
    "medium.en": "openai/whisper-medium.en",
    "medium": "openai/whisper-medium",
    "large-v1": "openai/whisper-large-v1",
    "large-v2": "openai/whisper-large-v2",
    "large-v3": "openai/whisper-large-v3",
    "large": "openai/whisper-large-v3",
}


def get_device(device: str = None) -> str:
    if device:
        return device
    if torch.cuda.is_available():
        return 'cuda:0'
    if (mps := getattr(torch.backends, 'mps', None)) is not None and mps.is_available():
        return 'mps'
    return 'cpu'


def load_hf_pipe(model_name: str, device: str = None, flash: bool = False):
    from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
    device = get_device(device)
    is_cpu = (device if isinstance(device, str) else getattr(device, 'type', None)) == 'cpu'
    dtype = torch.float32 if is_cpu or not torch.cuda.is_available() else torch.float16
    model_id = HF_MODELS.get(model_name, model_name)
    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        model_id,
        torch_dtype=dtype,
        low_cpu_mem_usage=True,
        use_safetensors=True,
        use_flash_attention_2=flash
    ).to(device)

    processor = AutoProcessor.from_pretrained(model_id)

    if not flash:
        try:
            model = model.to_bettertransformer()
        except ValueError:
            pass

    pipe = pipeline(
        "automatic-speech-recognition",
        model=model,
        tokenizer=processor.tokenizer,
        feature_extractor=processor.feature_extractor,
        max_new_tokens=128,
        chunk_length_s=30,
        torch_dtype=dtype,
        device=device,
    )

    return pipe


class GenericWrapper:

    def __init__(self, obj, prompt_ids=None):
        self.obj = obj
        self.prompt_ids = None if prompt_ids is None or len(prompt_ids) == 0 else prompt_ids

    def __getattr__(self, attr):
        val = getattr(self.obj, attr)
        return val


class InnerIteratorWrapper(GenericWrapper):

    def __next__(self):
        output = next(self.obj)
        if self.prompt_ids is None:
            return output
        prompt_len = len(self.prompt_ids)
        for chunk in output:
            if (chunk['tokens'][..., :prompt_len] == self.prompt_ids).all():
                chunk['tokens'] = chunk['tokens'][..., prompt_len:]
            if 'token_timestamps' in chunk and (chunk['token_timestamps'][..., :prompt_len] == self.prompt_ids).all():
                chunk['token_timestamps'] = chunk['token_timestamps'][..., prompt_len:]
        return output


class IteratorWrapper(GenericWrapper):

    def __iter__(self):
        iter(self.obj)
        self.obj.iterator = InnerIteratorWrapper(self.obj.iterator, self.prompt_ids)
        return self.obj


class PipelineWrapper(GenericWrapper):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.super_mode = False
        self.call_kwargs = {}

    def __getattr__(self, attr):
        non_parent_val = super().__getattr__(attr)
        if not self.super_mode:
            return non_parent_val
        val = getattr(super(type(self.obj), self.obj), attr, non_parent_val)
        if hasattr(val, '__isabstractmethod__') and val.__isabstractmethod__:
            val = non_parent_val
        if attr == 'get_iterator':

            def get_iterator(*_args, **_kwargs):
                iterator = val(*_args, **_kwargs)
                iterator = IteratorWrapper(iterator, self.prompt_ids)
                return iterator

            return get_iterator

        elif attr == 'run_single':

            def run_single(*_args, **_kwargs):
                curr_class = self.obj.__class__
                if self.super_mode:
                    curr_class = curr_class.__base__
                if self.framework == 'pt' and curr_class.__name__ == 'ChunkPipeline':
                    _args = list(_args)
                    _args = (
                            [[_args[0]]] +
                            [
                                self.call_kwargs.get('num_workers') or self._num_workers or 0,
                                self.call_kwargs.get('batch_size') or self._batch_size or 1,
                            ] +
                            _args[1:]
                    )
                    output = next(iter(self.get_iterator(*_args, **_kwargs)))
                else:
                    output = val(*_args, **_kwargs)

                return output

            return run_single

        return val

    def __call__(self, *args, **kwargs):
        call = self.obj.__class__.__base__.__call__
        self.call_kwargs = kwargs
        self.super_mode = True
        return call(self, *args, **kwargs)


class WhisperHF:

    def __init__(self, model_name: str, device: str = None, flash: bool = False, pipeline=None):
        self._model_name = model_name
        self._pipe = load_hf_pipe(self._model_name, device, flash=flash) if pipeline is None else pipeline
        self._model_name = getattr(self._pipe.model, 'name_or_path', self._model_name)

    @property
    def sampling_rate(self):
        return self._pipe.feature_extractor.sampling_rate

    @property
    def model_name(self):
        if self._model_name is None:
            return getattr(self._pipe.model, 'name_or_path', 'n/a')
        return self._model_name

    def _modify_pipe_iterator(self, prompt_ids=None):
        if not isinstance(self._pipe, PipelineWrapper):
            self._pipe = PipelineWrapper(self._pipe, prompt_ids)

    def _inner_transcribe(
            self,
            audio: Union[str, bytes, np.ndarray],
            language: Optional[str] = None,
            task: Optional[str] = 'transcribe',
            batch_size: Optional[int] = 24,
            word_timestamps=True,
            verbose: Optional[bool] = False,
            initial_prompt: Optional[str] = None,
            **kwargs
    ):
        generate_kwargs = {}
        if self.model_name.endswith('en'):
            language = task = None
        if task is not None:
            generate_kwargs['task'] = task
        if language is not None:
            generate_kwargs['language'] = language
        prompt_ids = kwargs.pop('prompt_ids', None)
        if initial_prompt:
            if prompt_ids:
                warnings.warn(
                    'Got both ``prompt_ids`` and ``initial_prompt`` but ``prompt_ids`` will be ignored.',
                    stacklevel=2
                )
            prompt_ids = self._pipe.tokenizer.get_prompt_ids(initial_prompt, 'pt')
        if prompt_ids is not None:
            self._modify_pipe_iterator(prompt_ids)
            generate_kwargs['prompt_ids'] = prompt_ids
        generate_kwargs.update(kwargs)
        if verbose is not None:
            print(f'Transcribing with Hugging Face Whisper ({self.model_name})...')
        pipe_kwargs = dict(
            generate_kwargs=generate_kwargs,
            return_timestamps='word' if word_timestamps else True
        )
        if batch_size is not None:
            pipe_kwargs['batch_size'] = batch_size
        output = self._pipe(audio, **pipe_kwargs)
        result = output['chunks']
        if verbose is not None:
            print(f'Transcription completed.')

        def replace_none_ts(parts):
            total_dur = round(audio.shape[-1] / self.sampling_rate, 3) if isinstance(audio, np.ndarray) else None
            _medium_dur = _ts_nonzero_mask = None

            def ts_nonzero_mask() -> np.ndarray:
                nonlocal _ts_nonzero_mask
                if _ts_nonzero_mask is None:
                    _ts_nonzero_mask = np.array([(p['end'] or p['start']) is not None for p in parts])
                return _ts_nonzero_mask

            def medium_dur() -> float:
                nonlocal _medium_dur
                if _medium_dur is None:
                    nonzero_dus = [p['end'] - p['start'] for p in parts if None not in (p['end'], p['start'])]
                    nonzero_durs = np.array(nonzero_dus)
                    _medium_dur = np.median(nonzero_durs) * 2 if len(nonzero_durs) else 2.0
                return _medium_dur

            def _curr_max_end(start: float, next_idx: float) -> float:
                max_end = total_dur
                if next_idx != len(parts):
                    mask = np.flatnonzero(ts_nonzero_mask()[next_idx:])
                    if len(mask):
                        _part = parts[mask[0]+next_idx]
                        max_end = _part['start'] or _part['end']

                new_end = round(start + medium_dur(), 3)
                if max_end is None:
                    return new_end
                if new_end > max_end:
                    return max_end
                return new_end

            for i, part in enumerate(parts, 1):
                if part['start'] is None:
                    is_first = i == 1
                    if is_first:
                        new_start = round((part['end'] or 0) - medium_dur(), 3)
                        part['start'] = max(new_start, 0.0)
                    else:
                        part['start'] = parts[i - 2]['end']
                if part['end'] is None:
                    no_next_start = i == len(parts) or parts[i]['start'] is None
                    part['end'] = _curr_max_end(part['start'], i) if no_next_start else parts[i]['start']

        if word_timestamps:
            words = [
                dict(start=word['timestamp'][0], end=word['timestamp'][1], word=word['text'])
                for word in result
            ]
            replace_none_ts(words)
            return [words]
        segs = [
            dict(start=seg['timestamp'][0], end=seg['timestamp'][1], text=seg['text'])
            for seg in result
        ]
        replace_none_ts(segs)
        return segs

    def transcribe(
            self,
            audio: Union[str, bytes, np.ndarray],
            *,
            language: Optional[str] = None,
            task: Optional[str] = 'transcribe',
            batch_size: Optional[int] = 24,
            word_timestamps: bool = True,
            verbose: Optional[bool] = False,
            regroup: Union[bool, str] = True,
            suppress_silence: bool = True,
            suppress_word_ts: bool = True,
            use_word_position: bool = True,
            q_levels: int = 20,
            k_size: int = 5,
            denoiser: Optional[str] = None,
            denoiser_options: Optional[dict] = None,
            vad: bool = False,
            vad_threshold: float = 0.35,
            vad_onnx: bool = False,
            min_word_dur: Optional[float] = None,
            nonspeech_error: float = 0.1,
            only_voice_freq: bool = False,
            only_ffmpeg: bool = False,
            check_sorted: bool = True,
            **options
    ):
        transcribe_any_options = isolate_useful_options(options, transcribe_any, pop=True)
        denoiser, denoiser_options = convert_demucs_kwargs(
            denoiser, denoiser_options,
            demucs=transcribe_any_options.pop('demucs', None),
            demucs_options=transcribe_any_options.pop('demucs_options', None)
        )

        if isinstance(audio, (str, bytes)):
            audio = prep_audio(audio, sr=self.sampling_rate).numpy()
            transcribe_any_options['input_sr'] = self.sampling_rate

        if 'input_sr' not in transcribe_any_options:
            transcribe_any_options['input_sr'] = self.sampling_rate

        if denoiser or only_voice_freq:
            if 'audio_type' not in transcribe_any_options:
                transcribe_any_options['audio_type'] = 'numpy'
            if 'model_sr' not in transcribe_any_options:
                transcribe_any_options['model_sr'] = self.sampling_rate

        inference_kwargs = dict(
            audio=audio,
            language=language,
            task=task,
            batch_size=batch_size,
            word_timestamps=word_timestamps,
            verbose=verbose,
            **options
        )
        return transcribe_any(
            inference_func=self._inner_transcribe,
            audio=audio,
            inference_kwargs=inference_kwargs,
            verbose=verbose,
            regroup=regroup,
            suppress_silence=suppress_silence,
            suppress_word_ts=suppress_word_ts,
            q_levels=q_levels,
            k_size=k_size,
            denoiser=denoiser,
            denoiser_options=denoiser_options,
            vad=vad,
            vad_threshold=vad_threshold,
            vad_onnx=vad_onnx,
            min_word_dur=min_word_dur,
            nonspeech_error=nonspeech_error,
            use_word_position=use_word_position,
            only_voice_freq=only_voice_freq,
            only_ffmpeg=only_ffmpeg,
            force_order=True,
            check_sorted=check_sorted,
            **transcribe_any_options
        )


def load_hf_whisper(model_name: str, device: str = None, flash: bool = False, pipeline=None):
    return WhisperHF(model_name, device, flash=flash, pipeline=pipeline)
