from typing import Dict, Optional, Tuple

import mlx.core as mx
import mlx.nn as nn
from mlx_lm.models.cache import make_prompt_cache
from mlx_lm.models.llama import Model as LlamaModel
from mlx_lm.models.llama import ModelArgs as LlamaModelConfig
from mlx_lm.sample_utils import make_logits_processors, make_sampler

from ..config import LLAMA_CONFIGS, T3Config
from .cond_enc import T3Cond, T3CondEnc
from .learned_pos_emb import LearnedPositionEmbeddings


class T3(nn.Module):
    """
    Token-To-Token (T3) TTS model using LLaMA as backbone.

    Generates speech tokens from text tokens, conditioned on speaker embeddings
    and optional emotion/prompt controls.
    """

    def __init__(self, hp: Optional[T3Config] = None):
        super().__init__()

        if hp is None:
            hp = T3Config.english_only()

        self.hp = hp

        # Create LLaMA config from our T3 config
        llama_config_dict = LLAMA_CONFIGS[hp.llama_config_name].copy()
        self.cfg = LlamaModelConfig(**llama_config_dict)

        # LLaMA transformer backbone
        self.tfmr = LlamaModel(self.cfg)
        self.dim = self.cfg.hidden_size

        # Conditioning encoder
        self.cond_enc = T3CondEnc(hp)

        # Text and speech token embeddings
        self.text_emb = nn.Embedding(hp.text_tokens_dict_size, self.dim)
        self.speech_emb = nn.Embedding(hp.speech_tokens_dict_size, self.dim)

        # Learned position embeddings (optional)
        if hp.input_pos_emb == "learned":
            max_text_seq_len = hp.max_text_tokens + 2
            self.text_pos_emb = LearnedPositionEmbeddings(max_text_seq_len, self.dim)

            max_mel_seq_len = hp.max_speech_tokens + 2 + 2
            self.speech_pos_emb = LearnedPositionEmbeddings(max_mel_seq_len, self.dim)

        # Output projection heads
        self.text_head = nn.Linear(
            self.cfg.hidden_size, hp.text_tokens_dict_size, bias=False
        )
        self.speech_head = nn.Linear(
            self.cfg.hidden_size, hp.speech_tokens_dict_size, bias=False
        )

    def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]:
        """
        Sanitize PyTorch weights for MLX.

        T3 uses a LLaMA backbone which mlx_lm already handles.
        We need to handle:
        - Mapping tfmr.* weights to the LLaMA model format (tfmr.X -> tfmr.model.X)
        - Conv1d weight transposition in conditioning encoder

        Note: Perceiver weights use the same naming convention (to_q, to_k, to_v, proj_out)
        as the original PyTorch implementation, so no renaming is needed.

        This method is idempotent - it checks shapes before transposing to support
        both PyTorch-format and pre-converted MLX-format weights.
        """
        import re

        from mlx.utils import tree_flatten

        new_weights = {}

        # Get expected shapes from model for idempotent transposition
        curr_weights = dict(tree_flatten(self.parameters()))

        for key, value in weights.items():
            new_key = key

            # === Transformer weight name mapping ===
            # PyTorch uses: tfmr.layers.X, tfmr.embed_tokens, tfmr.norm
            # mlx_lm LlamaModel uses: model.layers.X, model.embed_tokens, model.norm
            # So we need: tfmr.layers.X -> tfmr.model.layers.X, etc.
            # Check if already converted (idempotent)
            if key.startswith("tfmr.") and not key.startswith("tfmr.model."):
                # Map tfmr.X to tfmr.model.X for transformer internal components
                # These are: layers, embed_tokens, norm
                patterns_to_prefix = [
                    r"^tfmr\.layers\.",
                    r"^tfmr\.embed_tokens\.",
                    r"^tfmr\.norm\.",
                ]
                for pattern in patterns_to_prefix:
                    if re.match(pattern, key):
                        new_key = re.sub(r"^tfmr\.", "tfmr.model.", key)
                        break

            # Conv1d weight transposition (idempotent)
            # Only transpose if shape doesn't match expected MLX format
            # PyTorch Conv1d: (out_channels, in_channels, kernel_size)
            # MLX Conv1d: (out_channels, kernel_size, in_channels)
            if "conv" in new_key.lower() and "weight" in new_key and value.ndim == 3:
                if (
                    new_key in curr_weights
                    and value.shape != curr_weights[new_key].shape
                ):
                    value = value.swapaxes(1, 2)

            new_weights[new_key] = value

        # Delegate to LLaMA sanitizer for transformer weights if it has one
        if hasattr(self.tfmr, "sanitize"):
            tfmr_weights = {
                k: v for k, v in new_weights.items() if k.startswith("tfmr.")
            }
            other_weights = {
                k: v for k, v in new_weights.items() if not k.startswith("tfmr.")
            }
            tfmr_weights = self.tfmr.sanitize(tfmr_weights)
            new_weights = {**other_weights, **tfmr_weights}

        return new_weights

    def prepare_conditioning(self, t3_cond: T3Cond) -> mx.array:
        """
        Prepare conditioning embeddings from T3Cond.

        Args:
            t3_cond: Conditioning information

        Returns:
            Conditioning embeddings (B, cond_len, dim)
        """
        # Embed speech prompt tokens if provided
        if (
            t3_cond.cond_prompt_speech_tokens is not None
            and t3_cond.cond_prompt_speech_emb is None
        ):
            t3_cond.cond_prompt_speech_emb = self.speech_emb(
                t3_cond.cond_prompt_speech_tokens
            ) + self.speech_pos_emb(t3_cond.cond_prompt_speech_tokens)

        return self.cond_enc(t3_cond)

    def prepare_input_embeds(
        self,
        t3_cond: T3Cond,
        text_tokens: mx.array,
        speech_tokens: mx.array,
        cfg_weight: float = 0.0,
    ) -> Tuple[mx.array, int]:
        """
        Prepare input embeddings for the transformer.

        Args:
            t3_cond: Conditioning information
            text_tokens: Text token IDs (B, text_len)
            speech_tokens: Speech token IDs (B, speech_len)
            cfg_weight: Classifier-free guidance weight

        Returns:
            Tuple of (embeddings, conditioning_length)
        """
        # Prepare conditioning embeddings
        cond_emb = self.prepare_conditioning(t3_cond)  # (B, len_cond, dim)

        # Text embeddings
        text_emb = self.text_emb(text_tokens)  # (B, len_text, dim)

        # CFG: zero out second batch item for unconditional
        if cfg_weight > 0.0 and text_emb.shape[0] > 1:
            text_emb = mx.concatenate(
                [
                    text_emb[:1],
                    mx.zeros_like(text_emb[1:2]),
                ],
                axis=0,
            )

        # Speech embeddings
        speech_emb = self.speech_emb(speech_tokens)  # (B, len_speech, dim)

        # Add position embeddings if using learned positions
        if self.hp.input_pos_emb == "learned":
            text_emb = text_emb + self.text_pos_emb(text_tokens)
            speech_emb = speech_emb + self.speech_pos_emb(speech_tokens)

        len_cond = cond_emb.shape[1]

        # Broadcast conditioning if batch sizes don't match
        if cond_emb.shape[0] != text_emb.shape[0]:
            cond_emb = mx.broadcast_to(
                cond_emb, (text_emb.shape[0],) + cond_emb.shape[1:]
            )

        # Broadcast speech embeddings if batch sizes don't match (e.g., CFG)
        if speech_emb.shape[0] != text_emb.shape[0]:
            speech_emb = mx.broadcast_to(
                speech_emb, (text_emb.shape[0],) + speech_emb.shape[1:]
            )

        # Concatenate: [conditioning | text | speech]
        embeds = mx.concatenate([cond_emb, text_emb, speech_emb], axis=1)

        return embeds, len_cond

    def __call__(
        self,
        t3_cond: T3Cond,
        text_tokens: mx.array,
        text_token_lens: mx.array,
        speech_tokens: mx.array,
        speech_token_lens: mx.array,
        cache=None,
    ) -> dict:
        """
        Forward pass through T3 model.

        Args:
            t3_cond: Conditioning information
            text_tokens: Text token IDs (B, text_len)
            text_token_lens: Valid lengths for each text sequence (B,)
            speech_tokens: Speech token IDs (B, speech_len)
            speech_token_lens: Valid lengths for each speech sequence (B,)
            cache: Optional KV cache for inference

        Returns:
            Dictionary with text_logits, speech_logits, and hidden_states
        """
        # Prepare input embeddings
        embeds, len_cond = self.prepare_input_embeds(
            t3_cond=t3_cond,
            text_tokens=text_tokens,
            speech_tokens=speech_tokens,
        )

        # Forward through LLaMA backbone
        # Note: mlx_lm's LlamaModel takes input_embeddings differently
        hidden_states = self.tfmr.model(
            inputs=None,
            cache=cache,
            input_embeddings=embeds,
        )

        # Extract text and speech portions of hidden states
        B = text_tokens.shape[0]
        len_text = text_tokens.shape[1]
        len_speech = speech_tokens.shape[1]
        dim = hidden_states.shape[-1]

        # Allocate output tensors
        text_latents = mx.zeros((B, len_text, dim))
        speech_latents = mx.zeros((B, len_speech, dim))

        # Split hidden states by sequence position
        for i in range(B):
            ttl = int(text_token_lens[i])
            stl = int(speech_token_lens[i])

            text_start = len_cond
            text_end = len_cond + ttl

            speech_start = len_cond + len_text
            speech_end = speech_start + stl

            text_latents = (
                mx.concatenate(
                    [
                        text_latents[:i],
                        hidden_states[i : i + 1, text_start:text_end, :],
                        text_latents[i + 1 :],
                    ],
                    axis=0,
                )
                if i < B
                else text_latents
            )

            speech_latents = (
                mx.concatenate(
                    [
                        speech_latents[:i],
                        hidden_states[i : i + 1, speech_start:speech_end, :],
                        speech_latents[i + 1 :],
                    ],
                    axis=0,
                )
                if i < B
                else speech_latents
            )

        # Project to vocabulary
        text_logits = self.text_head(text_latents)
        speech_logits = self.speech_head(speech_latents)

        return {
            "text_logits": text_logits,
            "text_latents": text_latents,
            "speech_logits": speech_logits,
            "speech_latents": speech_latents,
            "hidden_states": hidden_states,
        }

    def inference(
        self,
        t3_cond: T3Cond,
        text_tokens: mx.array,
        initial_speech_tokens: Optional[mx.array] = None,
        max_new_tokens: int = 1024,
        temperature: float = 0.8,
        top_p: float = 0.95,
        min_p: float = 0.05,
        repetition_penalty: float = 1.2,
        cfg_weight: float = 0.5,
    ) -> mx.array:
        """
        Generate speech tokens from text tokens.

        This matches the original PyTorch implementation which uses:
        - KV caching for efficient generation
        - Position embeddings for each generated token
        - Repetition penalty, min_p, and top_p filtering
        - Classifier-free guidance (CFG)

        Args:
            t3_cond: Conditioning information
            text_tokens: Text token IDs (1D or 2D)
            initial_speech_tokens: Optional initial speech tokens
            max_new_tokens: Maximum tokens to generate
            temperature: Sampling temperature
            top_p: Top-p sampling threshold
            min_p: Minimum probability threshold
            repetition_penalty: Repetition penalty factor
            cfg_weight: Classifier-free guidance weight

        Returns:
            Generated speech tokens (B, T)
        """
        # Ensure text_tokens is 2D
        if text_tokens.ndim == 1:
            text_tokens = mx.expand_dims(text_tokens, 0)

        # Default initial speech token (BOS)
        bos_token = mx.array([[self.hp.start_speech_token]], dtype=mx.int32)

        # Prepare conditioning and text embeddings
        embeds, len_cond = self.prepare_input_embeds(
            t3_cond=t3_cond,
            text_tokens=text_tokens,
            speech_tokens=bos_token,
            cfg_weight=cfg_weight,
        )

        # Create BOS embedding with position embedding at position 0
        bos_embed = self.speech_emb(bos_token)  # (1, 1, dim)
        bos_embed = bos_embed + self.speech_pos_emb.get_fixed_embedding(0)

        # For CFG, duplicate BOS embed
        if cfg_weight > 0.0:
            bos_embed = mx.concatenate([bos_embed, bos_embed], axis=0)

        # Combine conditioning+text embeddings with BOS token
        # embeds already has [cond | text | bos] structure from prepare_input_embeds
        # We need to replace the speech part with our position-embedded BOS
        # Actually, let's rebuild the input properly
        cond_emb = self.prepare_conditioning(t3_cond)  # (1, len_cond, dim)
        text_emb = self.text_emb(text_tokens)  # (B, len_text, dim)

        if cfg_weight > 0.0:
            # Zero out second batch for unconditional
            text_emb = mx.concatenate(
                [
                    text_emb[:1],
                    mx.zeros_like(text_emb[:1]),
                ],
                axis=0,
            )

        if self.hp.input_pos_emb == "learned":
            text_emb = text_emb + self.text_pos_emb(text_tokens)

        # Broadcast conditioning if needed
        if cond_emb.shape[0] != text_emb.shape[0]:
            cond_emb = mx.broadcast_to(
                cond_emb, (text_emb.shape[0],) + cond_emb.shape[1:]
            )

        # Build initial input: [cond | text | bos]
        input_embeddings = mx.concatenate([cond_emb, text_emb, bos_embed], axis=1)

        # Create KV cache
        cache = make_prompt_cache(self.tfmr)

        # Create sampler and logits processors
        sampler = make_sampler(temp=temperature, top_p=top_p, min_p=min_p)
        logits_processors = make_logits_processors(
            logit_bias=None,
            repetition_penalty=repetition_penalty,
            repetition_context_size=max_new_tokens,  # Use all generated tokens for repetition
        )

        # Track generated tokens (Python list for efficiency - avoid growing mx.array each step)
        generated_ids = [self.hp.start_speech_token]

        # Initial forward pass to fill cache
        hidden = self.tfmr.model(
            inputs=None, input_embeddings=input_embeddings, cache=cache
        )

        # Generation loop
        for step in range(max_new_tokens):
            # Get logits for last position
            logits = self.speech_head(hidden[:, -1:, :])  # (B, 1, vocab)
            logits = logits.squeeze(1)  # (B, vocab)

            # Apply CFG
            if cfg_weight > 0.0 and logits.shape[0] > 1:
                cond_logits = logits[0:1, :]
                uncond_logits = logits[1:2, :]
                logits = cond_logits + cfg_weight * (cond_logits - uncond_logits)
            else:
                logits = logits[0:1, :]

            # Apply logits processors (repetition penalty)
            # Lazily convert Python list to array - processor only uses last N tokens anyway
            for processor in logits_processors:
                tokens_for_penalty = mx.array([generated_ids], dtype=mx.int32)
                logits = processor(tokens_for_penalty, logits)

            # Sample next token using the sampler (handles temperature, top_p, min_p)
            next_token = sampler(logits)

            mx.eval(next_token)
            next_token_id = int(next_token[0])

            # Check for EOS
            if next_token_id == self.hp.stop_speech_token:
                generated_ids.append(next_token_id)
                break

            generated_ids.append(next_token_id)

            # Create embedding for next token with position embedding
            next_token_embed = self.speech_emb(mx.array([[next_token_id]]))
            next_token_embed = (
                next_token_embed + self.speech_pos_emb.get_fixed_embedding(step + 1)
            )

            # For CFG, duplicate
            if cfg_weight > 0.0:
                next_token_embed = mx.concatenate(
                    [next_token_embed, next_token_embed], axis=0
                )

            # Forward pass with cache (only new token)
            hidden = self.tfmr.model(
                inputs=None, input_embeddings=next_token_embed, cache=cache
            )

            mx.eval(hidden)

        return mx.array([generated_ids])
