from dataclasses import dataclass
from typing import Dict
from typing import Iterable, Optional
import time
import numpy as np
import torch
import torch.nn.functional as F
from torch import Tensor
from torch import nn

import whisper

# import whisper_timestamped as whisper

from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank

from funasr.register import tables


@tables.register("model_classes", "Whisper-tiny.en")
@tables.register("model_classes", "Whisper-tiny")
@tables.register("model_classes", "Whisper-base.en")
@tables.register("model_classes", "Whisper-base")
@tables.register("model_classes", "Whisper-small.en")
@tables.register("model_classes", "Whisper-small")
@tables.register("model_classes", "Whisper-medium.en")
@tables.register("model_classes", "Whisper-medium")
@tables.register("model_classes", "Whisper-large-v1")
@tables.register("model_classes", "Whisper-large-v2")
@tables.register("model_classes", "Whisper-large-v3")
@tables.register("model_classes", "Whisper-large-v3-turbo")
@tables.register("model_classes", "WhisperWarp")
class WhisperWarp(nn.Module):
    """Whisper: OpenAI Whisper model integration.

    Wraps Whisper for multilingual speech recognition and translation
    within FunASR's AutoModel interface.

    Supports: whisper-tiny through whisper-large-v3-turbo.
    Output: {"key": str, "text": str}
    """

    def __init__(self, *args, **kwargs):
        """Initialize WhisperWarp.
        
            Args:
                *args: Variable positional arguments.
                **kwargs: Additional keyword arguments.
            """
        super().__init__()
        hub = kwargs.get("hub", "funasr")
        if hub == "openai":
            model_or_path = kwargs.get("model_path", "Whisper-large-v3")
            if model_or_path.startswith("Whisper-"):
                model_or_path = model_or_path.replace("Whisper-", "")
            model = whisper.load_model(model_or_path)
        else:
            dims = kwargs.get("dims", {})
            dims = whisper.model.ModelDimensions(**dims)
            model = whisper.model.Whisper(dims=dims)

        self.model = model

        self.encoder_output_size = self.model.dims.n_audio_state

    def forward(
        self,
        speech: torch.Tensor = None,
        speech_lengths: torch.Tensor = None,
        text: torch.Tensor = None,
        text_lengths: torch.Tensor = None,
        **kwargs,
    ):
        """Forward pass for training. Computes cross-entropy loss.

        Args:
            speech: (B, T, D) mel-spectrogram features
            speech_lengths: (B,) lengths of each audio
            text: (B, U) token IDs (with SOT/EOT tokens)
            text_lengths: (B,) lengths of each text sequence
        Returns:
            dict with "loss" and optionally "stats"
        """
        if speech is None or text is None:
            raise ValueError("forward() requires speech and text for training")

        # Encoder
        audio_features = self.model.encoder(speech)

        # Decoder: shift text right for teacher forcing
        # text format: [SOT, lang, task, ..., tokens, EOT]
        decoder_input = text[:, :-1]
        decoder_target = text[:, 1:]

        # Decoder forward
        logits = self.model.decoder(decoder_input, audio_features)

        # Cross-entropy loss (ignore padding, token_id = -1 or pad)
        loss = F.cross_entropy(
            logits.reshape(-1, logits.size(-1)),
            decoder_target.reshape(-1),
            ignore_index=-100,
        )

        stats = {
            "loss": loss.detach().item(),
            "batch_size": speech.size(0),
        }
        return {"loss": loss, "stats": stats}

    def inference(
        self,
        data_in,
        data_lengths=None,
        key: list = None,
        tokenizer=None,
        frontend=None,
        **kwargs,
    ):
        """Run inference on input data.
        
            Args:
                data_in: Input data (audio samples, file paths, or text).
                data_lengths: Lengths of each input sample in the batch.
                key: Sample identifiers.
                tokenizer: Tokenizer instance for text encoding/decoding.
                frontend: Audio frontend for feature extraction.
                **kwargs: Additional keyword arguments.
            """
        if kwargs.get("batch_size", 1) > 1:
            raise NotImplementedError("batch decoding is not implemented")

        if frontend is None and not hasattr(self, "frontend"):
            frontend_class = tables.frontend_classes.get("WhisperFrontend")
            frontend = frontend_class(
                n_mels=self.model.dims.n_mels, do_pad_trim=kwargs.get("do_pad_trim", True)
            )
            self.frontend = frontend
        else:
            frontend = frontend if frontend is not None else self.frontend

        meta_data = {}
        if (
            isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank"
        ):  # fbank
            speech, speech_lengths = data_in, data_lengths
            if len(speech.shape) < 3:
                speech = speech[None, :, :]
            if speech_lengths is None:
                speech_lengths = speech.shape[1]
        else:
            # extract fbank feats
            time1 = time.perf_counter()
            audio_sample_list = load_audio_text_image_video(
                data_in,
                fs=frontend.fs if hasattr(frontend, "fs") else 16000,
                audio_fs=kwargs.get("fs", 16000),
                data_type=kwargs.get("data_type", "sound"),
                tokenizer=tokenizer,
            )
            time2 = time.perf_counter()
            meta_data["load_data"] = f"{time2 - time1:0.3f}"
            speech, speech_lengths = extract_fbank(
                audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend
            )
            time3 = time.perf_counter()
            meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
            frame_shift = frontend.frame_shift if hasattr(frontend, "frame_shift") else 10
            lfr_n = frontend.lfr_n if hasattr(frontend, "lfr_n") else 1
            meta_data["batch_data_time"] = speech_lengths.sum().item() * frame_shift * lfr_n / 1000

        speech = speech.to(device=kwargs["device"])[0, :, :]
        speech_lengths = speech_lengths.to(device=kwargs["device"])

        # # detect the spoken language
        # _, probs = self.model.detect_language(speech)
        # print(f"Detected language: {max(probs, key=probs.get)}")

        # decode the audio
        options = whisper.DecodingOptions(**kwargs.get("DecodingOptions", {}))

        result = whisper.decode(self.model, speech, options=options)
        # result = whisper.transcribe(self.model, speech)

        results = []
        result_i = {"key": key[0], "text": result.text}

        results.append(result_i)

        return results, meta_data