import glob
import math
import time
from typing import Any, Callable, Dict, Generator, List, Optional, Tuple

import mlx.core as mx
import mlx.nn as nn
import numpy as np

from mlx_audio.stt.generate import wired_limit
from mlx_audio.stt.utils import get_model_path

from ..base import STTOutput
from .config import AudioConfig, ModelConfig


class Attention(nn.Module):
    def __init__(
        self,
        config: ModelConfig,
    ):
        super().__init__()
        self.embed_dim = config.d_model
        self.num_heads = config.encoder_attention_heads
        self.dropout = config.attention_dropout
        self.head_dim = config.d_model // config.encoder_attention_heads
        self.config = config

        if (self.head_dim * self.num_heads) != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
                f" and `num_heads`: {self.num_heads})."
            )
        self.scaling = self.head_dim**-0.5

        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)

    def __call__(
        self,
        x: mx.array,
        mask: Optional[mx.array] = None,
    ) -> Tuple[mx.array, Optional[mx.array]]:
        bsz, tgt_len, _ = x.shape

        query_states = self.q_proj(x) * self.scaling
        key_states = self.k_proj(x)
        value_states = self.v_proj(x)

        query_states = query_states.reshape(
            bsz, tgt_len, self.num_heads, self.head_dim
        ).transpose(0, 2, 1, 3)
        key_states = key_states.reshape(
            bsz, -1, self.num_heads, self.head_dim
        ).transpose(0, 2, 1, 3)
        value_states = value_states.reshape(
            bsz, -1, self.num_heads, self.head_dim
        ).transpose(0, 2, 1, 3)

        attn_output = mx.fast.scaled_dot_product_attention(
            query_states, key_states, value_states, scale=1.0, mask=mask
        )

        attn_output = attn_output.transpose(0, 2, 1, 3).reshape(
            bsz, tgt_len, self.embed_dim
        )

        return self.out_proj(attn_output)


class VoxtralEncoderLayer(nn.Module):
    def __init__(self, config: AudioConfig):
        super().__init__()
        self.embed_dim = config.d_model

        self.self_attn = Attention(
            config=config,
        )
        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)

        self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
        self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
        self.final_layer_norm = nn.LayerNorm(self.embed_dim)

    def __call__(
        self,
        x: mx.array,
        mask: Optional[mx.array] = None,
    ) -> mx.array:
        r = x
        x = self.self_attn_layer_norm(x)
        x = self.self_attn(x, mask=mask)
        x = r + x

        r = x
        x = self.final_layer_norm(x)
        x = nn.gelu(self.fc1(x))
        x = self.fc2(x)
        x = r + x

        return x


class Encoder(nn.Module):
    def __init__(self, config: AudioConfig):
        super().__init__()
        self.config = config
        self.dropout = config.dropout
        self.layerdrop = config.encoder_layerdrop

        embed_dim = config.d_model
        self.num_mel_bins = config.num_mel_bins
        self.max_source_positions = config.max_source_positions
        self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0

        self.conv1 = nn.Conv1d(self.num_mel_bins, embed_dim, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1)

        self.embed_positions = nn.Embedding(self.max_source_positions, embed_dim)
        self.layers = [
            VoxtralEncoderLayer(config) for _ in range(config.encoder_layers)
        ]
        self.layer_norm = nn.LayerNorm(config.d_model)

    def __call__(
        self,
        x: mx.array,
        mask: Optional[mx.array] = None,
    ) -> mx.array:

        x = nn.gelu(self.conv1(x))
        x = nn.gelu(self.conv2(x))

        embed_pos = self.embed_positions.weight

        x = (x + embed_pos).astype(x.dtype)

        for encoder_layer in self.layers:
            x = encoder_layer(x, mask)

        return self.layer_norm(x)


class MultiModalProjector(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.linear_1 = nn.Linear(
            config.audio_config.intermediate_size,
            config.text_config.hidden_size,
            bias=False,
        )
        self.linear_2 = nn.Linear(
            config.text_config.hidden_size, config.text_config.hidden_size, bias=False
        )

    def __call__(self, audio_features: mx.array) -> mx.array:
        hidden_states = self.linear_1(audio_features)
        hidden_states = nn.gelu(hidden_states)
        hidden_states = self.linear_2(hidden_states)
        return hidden_states


class LanguageModel(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        self.model_type = config.model_type

        from mlx_lm.models.llama import LlamaModel

        self.model = LlamaModel(config)

        if not config.tie_word_embeddings:
            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

    def __call__(
        self,
        inputs: Optional[mx.array] = None,
        cache: Optional[mx.array] = None,
        input_embeddings: Optional[mx.array] = None,
    ):
        out = self.model(inputs, cache=cache, input_embeddings=input_embeddings)
        if self.config.tie_word_embeddings:
            out = self.model.embed_tokens.as_linear(out)
        else:
            out = self.lm_head(out)
        return out

    @property
    def layers(self):
        return self.model.layers


class Model(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        self.vocab_size = config.text_config.vocab_size

        self.language_model = LanguageModel(config.text_config)

        self.audio_tower = Encoder(config.audio_config)
        self.multi_modal_projector = MultiModalProjector(config)

    def get_audio_embeds(self, x: mx.array) -> mx.array:
        audio_embeds = self.audio_tower(x).reshape(
            -1, self.config.audio_config.intermediate_size
        )
        audio_embeds = self.multi_modal_projector(audio_embeds)
        return audio_embeds

    def _merge_input_embeddings(
        self,
        input_ids: mx.array,
        input_features: mx.array,
        cache: Optional[mx.array] = None,
    ) -> mx.array:
        if input_ids is not None:
            inputs_embeds = self.language_model.model.embed_tokens(input_ids)
        else:
            inputs_embeds = None

        if input_features is not None and (cache is None or cache[0].offset == 0):
            audio_embeds = self.get_audio_embeds(input_features)

            if inputs_embeds is not None:
                # Replace audio token placeholders with audio embeddings
                audio_token_mask = input_ids == self.config.audio_token_id
                # Expand audio_token_mask to match inputs_embeds shape
                audio_token_mask_expanded = audio_token_mask[..., None]
                audio_token_mask_expanded = mx.broadcast_to(
                    audio_token_mask_expanded, inputs_embeds.shape
                )
                # Expand audio_embeds to match the number of audio tokens
                audio_token_positions = np.where(audio_token_mask.flatten())[0].tolist()
                inputs_embeds_flat = inputs_embeds.reshape(-1, inputs_embeds.shape[-1])
                inputs_embeds_flat[audio_token_positions] = audio_embeds
                inputs_embeds = inputs_embeds_flat.reshape(inputs_embeds.shape)
            else:
                inputs_embeds = audio_embeds

        return inputs_embeds

    def __call__(
        self,
        input_ids: mx.array,
        input_features: mx.array = None,
        cache: Optional[mx.array] = None,
    ) -> mx.array:

        inputs_embeds = self._merge_input_embeddings(
            input_ids=input_ids,
            input_features=input_features,
            cache=cache,
        )

        logits = self.language_model(input_embeddings=inputs_embeds, cache=cache)

        return logits

    def sanitize(self, weights):
        sanitized_weights = {}
        for k, v in weights.items():
            if "conv" in k and "weight" in k:
                if v.shape[-1] < v.shape[-2]:
                    sanitized_weights[k] = v.transpose(0, 2, 1)
                else:
                    sanitized_weights[k] = v
            else:
                sanitized_weights[k] = v
        return sanitized_weights

    @classmethod
    def from_pretrained(
        cls,
        model_path: str,
        config: Optional[ModelConfig] = None,
        **kwargs,
    ):
        from transformers import AutoProcessor

        processor = AutoProcessor.from_pretrained(model_path)
        model_repo = model_path
        revision = kwargs.get("revision", None)
        force_download = kwargs.get("force_download", False)
        model_path = get_model_path(
            model_path, revision=revision, force_download=force_download
        )

        if config is None:
            import json

            with open(f"{model_path}/config.json", "r") as f:
                config_dict = json.load(f)
            config = ModelConfig.from_dict(config_dict)

        model = cls(config)
        model._processor = processor
        model._processor.tokenizer.eos_token_ids = getattr(
            model._processor.tokenizer, "eos_token_ids", [2, 4, 32000]
        )
        model.config.model_repo = model_repo

        weights = {}
        weight_files = glob.glob(str(model_path / "model-*.safetensors"))
        for file in weight_files:
            weights.update(mx.load(file))

        weights = model.sanitize(weights)

        model.load_weights(list(weights.items()))

        return model

    def stream_generate(
        self,
        input_ids: Optional[mx.array] = None,
        *,
        input_features: Optional[mx.array] = None,
        max_tokens: int = 128,
        sampler: Optional[Callable[mx.array, mx.array]] = None,
        generation_stream: bool = False,
    ) -> Generator[Tuple[mx.array, mx.array], None, None]:

        from mlx_lm.generate import generate_step

        input_embeddings = self._merge_input_embeddings(
            input_ids=input_ids,
            input_features=input_features,
        )[0]

        with wired_limit(self, [generation_stream]):
            for n, (token, logprobs) in enumerate(
                generate_step(
                    prompt=mx.array([]),
                    input_embeddings=input_embeddings,
                    model=self.language_model,
                    max_tokens=max_tokens,
                    sampler=sampler,
                )
            ):
                if token in self._processor.tokenizer.eos_token_ids:
                    break

                yield token, logprobs

    def generate(
        self,
        audio: List[mx.array],
        *,
        message: Optional[List[Dict[str, Any]]] = None,
        max_tokens: int = 128,
        temperature: float = 0.0,
        top_p: float = 0.95,
        top_k: int = 0,
        min_p: float = 0.0,
        min_tokens_to_keep: int = 1,
        language: str = "en",
        verbose: bool = False,
        generation_stream: bool = False,
    ) -> mx.array:

        start_time = time.time()

        if message is None:
            messages = [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "audio",
                            "path": audio,
                        },
                    ],
                }
            ]

        inputs = self._processor.apply_transcription_request(
            language=language, audio=audio, model_id=self.config.model_repo
        )
        input_ids = mx.array(inputs["input_ids"])
        input_features = mx.array(inputs["input_features"]).transpose(0, 2, 1)

        generated = []

        from mlx_lm.sample_utils import make_sampler

        sampler = make_sampler(
            temperature,
            top_p,
            min_p,
            min_tokens_to_keep=min_tokens_to_keep,
            top_k=top_k,
            xtc_special_tokens=self._processor.tokenizer.encode("\n")
            + list(self._processor.tokenizer.eos_token_ids),
        )

        for token, _ in self.stream_generate(
            input_ids=input_ids,
            input_features=input_features,
            max_tokens=max_tokens,
            sampler=sampler,
            generation_stream=generation_stream,
        ):
            generated.append(token)
            if verbose:
                print(self._processor.decode([token]), end="", flush=True)

        end_time = time.time()

        # Clear cache after each segment to avoid memory leaks
        mx.clear_cache()

        return STTOutput(
            text=self._processor.decode(generated),
            prompt_tokens=input_ids.shape[1],
            generation_tokens=len(generated),
            total_tokens=input_ids.shape[1] + len(generated),
            total_time=end_time - start_time,
            prompt_tps=input_ids.shape[1] / (end_time - start_time),
            generation_tps=len(generated) / (end_time - start_time),
        )
