"""Voxtral Realtime 4B - Main model orchestrator.

Inference pipeline:
1. Resample audio to 16kHz, pad (left silence + right silence)
2. Compute mel spectrogram
3. Run causal encoder -> 4x downsample -> adapter
4. Construct prompt: [BOS] + [STREAMING_PAD] * (n_left_pad + n_delay)
5. For each position: input = audio_embed + tok_embed(token_id)
6. Prefill decoder, then autoregressive generation until EOS
7. Decode tokens via Tekken tokenizer

Optimizations:
- Explicit mx.eval after prefill and periodically during decode to bound graph size
- Sampling uses logits directly (skips softmax+log round-trip)
"""

import time
from pathlib import Path
from typing import List, Optional, Union

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

from ..base import STTOutput
from .audio import compute_mel_filters, compute_mel_spectrogram
from .config import (
    RAW_AUDIO_LENGTH_PER_TOK,
    SAMPLE_RATE,
    ModelConfig,
    _num_delay_tokens,
)
from .decoder import Decoder, compute_time_embedding
from .encoder import AudioEncoder
from .tokenizer import TekkenTokenizer


def _pad_audio_streaming(audio_array, n_left_pad_tokens, n_right_pad_tokens):
    """Pad audio for offline streaming mode.

    Left pad: n_left_pad_tokens * 1280 samples of silence
    Right pad: align to 1280 + n_right_pad_tokens * 1280 of silence
    """
    mult_of = RAW_AUDIO_LENGTH_PER_TOK
    n_samples = len(audio_array)
    align_pad = (mult_of - (n_samples % mult_of)) % mult_of
    right_pad = align_pad + n_right_pad_tokens * mult_of
    left_pad = n_left_pad_tokens * mult_of
    return np.pad(audio_array, (left_pad, right_pad))


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

        self.encoder = AudioEncoder(config.encoder_args)
        self.decoder = Decoder(config.decoder)

        # Will be set in post_load_hook
        self._tokenizer = None
        self._mel_filters = None
        # Sentinel so _ensure_ada_scales is callable even if post_load_hook
        # has not run yet (e.g. unit tests that build Model directly).
        self._ada_scale_delay = -1

    def _ensure_mel_filters(self):
        if self._mel_filters is None:
            aec = self.config.audio_encoding_args
            filters_np = compute_mel_filters(
                num_mel_bins=aec.num_mel_bins,
                window_size=aec.window_size,
                sample_rate=aec.sampling_rate,
            )
            self._mel_filters = mx.array(filters_np, dtype=mx.float32)
        return self._mel_filters

    def _load_audio(
        self, audio_input: Union[str, Path, List[mx.array], mx.array, np.ndarray]
    ) -> np.ndarray:
        """Load and resample audio from file path or array to 16kHz float32 numpy."""
        if isinstance(audio_input, (str, Path)):
            from mlx_audio.audio_io import read as audio_read
            from mlx_audio.utils import resample_audio

            audio_np, sr = audio_read(str(audio_input), dtype="float32")
            audio_np = audio_np.flatten()
            if sr != SAMPLE_RATE:
                audio_np = resample_audio(audio_np, sr, SAMPLE_RATE)
            return audio_np
        if isinstance(audio_input, list):
            audio_input = audio_input[0]
        return np.array(audio_input).flatten().astype(np.float32)

    def _prepare_mel(self, audio_np, transcription_delay_ms=None):
        """Prepare mel spectrogram from audio numpy array."""
        delay_ms = transcription_delay_ms or self.config.transcription_delay_ms
        n_delay = _num_delay_tokens(delay_ms)
        n_left = self.config.n_left_pad_tokens
        n_right = (n_delay + 1) + 10

        padded = _pad_audio_streaming(audio_np, n_left, n_right)

        aec = self.config.audio_encoding_args
        mel_filters = self._ensure_mel_filters()
        audio_mx = mx.array(padded, dtype=mx.float32)
        mel = compute_mel_spectrogram(
            audio_mx,
            mel_filters,
            window_size=aec.window_size,
            hop_length=aec.hop_length,
            global_log_mel_max=aec.global_log_mel_max,
        )

        if mel.shape[1] % 2 != 0:
            mel = mel[:, 1:]

        return mel, n_delay

    def _ensure_ada_scales(self, transcription_delay_ms=None):
        """Ensure ada_scales match the given delay. Recomputes if needed."""
        delay_ms = transcription_delay_ms or self.config.transcription_delay_ms
        n_delay = _num_delay_tokens(delay_ms)
        if n_delay != self._ada_scale_delay:
            from .decoder import compute_time_embedding

            t_cond = compute_time_embedding(float(n_delay), self.config.decoder.dim)
            # Match the dtype the AdaRMSNorm weights were demoted to so the
            # matmul doesn't silently upcast (see post_load_hook).
            ada_weight = self.decoder.layers[0].ada_rms_norm_t_cond.ada_down.weight
            t_cond = t_cond.astype(ada_weight.dtype)
            self.decoder.precompute_ada_scales(t_cond)
            for scale in self.decoder._ada_scales:
                if scale is not None:
                    mx.eval(scale)
            self._ada_scale_delay = n_delay

    def _encode_and_prefill(self, audio_np, verbose=False, transcription_delay_ms=None):
        """Shared encoder + prefill logic with incremental encoding.

        Only encodes the first audio chunk before prefill, deferring the
        rest to the autoregressive loop. Returns a chunk generator for
        on-demand encoding of remaining audio.

        Returns:
            (adapter_out, n_audio_total, prompt_len, logits, cache,
             enc_chunk_gen, start_time)
        """
        start_time = time.time()

        self._ensure_ada_scales(transcription_delay_ms)
        mel, n_delay = self._prepare_mel(audio_np, transcription_delay_ms)

        # Run conv stem and compute total audio token count
        conv_out = self.encoder.conv_stem(mel)
        ds = self.encoder.config.downsample_factor
        n_audio_total = conv_out.shape[0] // ds

        n_left = self.config.n_left_pad_tokens
        prompt_len = 1 + n_left + n_delay
        sw = self.encoder.config.sliding_window

        if conv_out.shape[0] <= sw:
            # Short audio: use non-chunked path with optimized causal attention
            # (SDPA Flash Attention kernel, no RotatingKVCache overhead)
            adapter_out = self.encoder.encode_full(conv_out)
            enc_chunk_gen = None
        else:
            # Long audio: chunked encoding with generator for incremental decode
            enc_chunk_gen = self.encoder.encode_chunks(conv_out)
            adapter_chunks = []
            adapter_len = 0
            while adapter_len < prompt_len:
                try:
                    chunk_out = next(enc_chunk_gen)
                    chunk_adapter = self.encoder.downsample_and_project(chunk_out)
                    adapter_chunks.append(chunk_adapter)
                    adapter_len += chunk_adapter.shape[0]
                except StopIteration:
                    enc_chunk_gen = None
                    break
            adapter_out = (
                mx.concatenate(adapter_chunks, axis=0) if adapter_chunks else None
            )

        if verbose:
            print(f"Audio: {len(audio_np)} samples ({len(audio_np)/SAMPLE_RATE:.1f}s)")
            print(
                f"Encoder: {'non-chunked (causal)' if conv_out.shape[0] <= sw else 'chunked'}, "
                f"{adapter_out.shape[0]} adapter tokens"
            )
            print(f"Total audio tokens: {n_audio_total}")

        # Build prompt embeddings
        prompt_ids = [self.config.bos_token_id] + [
            self.config.streaming_pad_token_id
        ] * (n_left + n_delay)

        prompt_ids_mx = mx.array(prompt_ids)
        prompt_text_embeds = self.decoder.embed_tokens(prompt_ids_mx)
        prefix_embeds = adapter_out[:prompt_len] + prompt_text_embeds

        if verbose:
            print(f"Prompt: {prompt_len} tokens, Audio span: {n_audio_total} tokens")

        # Prefill (single fused forward pass)
        prefill_start = time.time()
        h, cache = self.decoder.forward(prefix_embeds, start_pos=0)
        logits = self.decoder.logits(h[-1])
        cache_arrays = [a for c in cache for a in (c.keys, c.values) if a is not None]
        mx.eval(logits, *cache_arrays)

        if verbose:
            total_ttft = time.time() - start_time
            prefill_time = time.time() - prefill_start
            print(
                f"Prefill: {prompt_len} tokens in {prefill_time*1000:.0f}ms "
                f"({prompt_len / prefill_time:.0f} tok/s)"
            )
            print(f"Total TTFT: {total_ttft*1000:.0f}ms")

        return (
            adapter_out,
            n_audio_total,
            prompt_len,
            logits,
            cache,
            enc_chunk_gen,
            start_time,
        )

    def generate(
        self,
        audio: Union[str, Path, List[mx.array], mx.array],
        *,
        max_tokens: int = 4096,
        temperature: float = 0.0,
        verbose: bool = False,
        stream: bool = False,
        transcription_delay_ms: Optional[int] = None,
        **kwargs,
    ):
        """Transcribe audio. Returns STTOutput, or yields text deltas if stream=True.

        Args:
            transcription_delay_ms: Override the model's default transcription delay
                (default 480ms). Lower values reduce latency but may hurt accuracy.
        """
        audio_np = self._load_audio(audio)

        if stream:
            return self._generate_stream(
                audio_np, max_tokens, temperature, verbose, transcription_delay_ms
            )

        adapter_out, n_audio, prompt_len, logits, cache, enc_chunk_gen, start_time = (
            self._encode_and_prefill(audio_np, verbose, transcription_delay_ms)
        )

        adapter_len = adapter_out.shape[0]
        generated = []

        # Double-buffered async eval: submit token computation to GPU,
        # read result at start of next iteration while GPU works on new ops
        next_tok = self._next_token_mx(logits, temperature)
        mx.async_eval(next_tok)

        decode_start = time.time()
        for pos in range(prompt_len, n_audio):
            token = int(next_tok.item())
            generated.append(token)
            if token == self.config.eos_token_id or len(generated) > max_tokens:
                break

            # Encode more audio chunks on demand
            if enc_chunk_gen is not None and pos >= adapter_len:
                try:
                    chunk_out = next(enc_chunk_gen)
                    chunk_adapter = self.encoder.downsample_and_project(chunk_out)
                    mx.eval(chunk_adapter)
                    adapter_out = mx.concatenate([adapter_out, chunk_adapter], axis=0)
                    adapter_len = adapter_out.shape[0]
                except StopIteration:
                    enc_chunk_gen = None

            if pos < adapter_len:
                embed = adapter_out[pos] + self.decoder.embed_token(token)
            else:
                embed = self.decoder.embed_token(token)

            h, cache = self.decoder.forward(embed[None, :], start_pos=pos, cache=cache)
            logits = self.decoder.logits(h.squeeze(0))
            next_tok = self._next_token_mx(logits, temperature)
            mx.async_eval(next_tok)

            if len(generated) % 256 == 0:
                mx.clear_cache()
        else:
            # Loop completed without break — read final pending token
            token = int(next_tok.item())
            generated.append(token)

        if generated and generated[-1] == self.config.eos_token_id:
            generated = generated[:-1]

        text = self._tokenizer.decode(generated).strip()
        end_time = time.time()
        total_time = end_time - start_time
        decode_time = end_time - decode_start

        if verbose:
            n_gen = len(generated)
            print(
                f"Decode: {n_gen} tokens in {decode_time:.3f}s "
                f"({n_gen / decode_time:.0f} tok/s, "
                f"{decode_time / max(n_gen, 1) * 1000:.1f} ms/tok)"
            )
            print(f"Total: {total_time:.3f}s")

        mx.clear_cache()

        return STTOutput(
            text=text,
            prompt_tokens=prompt_len,
            generation_tokens=len(generated),
            total_tokens=prompt_len + len(generated),
            total_time=total_time,
            prompt_tps=prompt_len / total_time if total_time > 0 else 0,
            generation_tps=len(generated) / decode_time if decode_time > 0 else 0,
        )

    def create_streaming_session(
        self,
        *,
        max_tokens: int = 4096,
        temperature: float = 0.0,
        transcription_delay_ms: Optional[int] = None,
    ):
        """Create a stateful streaming session for this model.

        Returns a ``VoxtralStreamingSession`` whose ``feed()`` is a cheap
        thread-safe queue push (safe to call from an asyncio loop), and
        whose ``step()`` / ``finalize_step()`` run one short unit of MLX
        work. Call ``step()`` repeatedly from an executor thread between
        audio feeds; between calls the executor is free for other work.
        """
        from .streaming import VoxtralStreamingSession

        return VoxtralStreamingSession(
            self,
            max_tokens=max_tokens,
            temperature=temperature,
            transcription_delay_ms=transcription_delay_ms,
        )

    def generate_streaming(
        self,
        source,
        *,
        max_tokens: int = 4096,
        temperature: float = 0.0,
        verbose: bool = False,
        transcription_delay_ms: Optional[int] = None,
    ):
        """Yield text deltas while audio is streamed in from ``source``.

        Thin wrapper around ``create_streaming_session`` that consumes the
        ``StreamingAudioSource`` in the calling thread (held for the life
        of the generation). For cooperative integrations that need to
        interleave MLX work with other tasks, prefer building your own
        loop on top of ``create_streaming_session`` instead.
        """
        sess = self.create_streaming_session(
            max_tokens=max_tokens,
            temperature=temperature,
            transcription_delay_ms=transcription_delay_ms,
        )

        start_time = time.time()
        while True:
            samples, closed = source.read()
            if samples.size > 0:
                sess.feed(samples)
            if closed:
                sess.close()
            for delta in sess.step(max_decode_tokens=16):
                yield delta
            if sess.done:
                break

        if verbose:
            total = time.time() - start_time
            print(
                f"Streaming generate: {len(sess.generated)} tokens in {total:.3f}s "
                f"({len(sess.generated) / total:.0f} tok/s)"
            )
        mx.clear_cache()

    def _generate_stream(
        self, audio_np, max_tokens, temperature, verbose, transcription_delay_ms=None
    ):
        """Generator that yields text deltas as tokens are decoded."""
        adapter_out, n_audio, prompt_len, logits, cache, enc_chunk_gen, start_time = (
            self._encode_and_prefill(audio_np, verbose, transcription_delay_ms)
        )

        adapter_len = adapter_out.shape[0]
        generated = []
        prev_text = ""

        next_tok = self._next_token_mx(logits, temperature)
        mx.async_eval(next_tok)

        for pos in range(prompt_len, n_audio):
            token = int(next_tok.item())
            generated.append(token)

            # Yield text delta (filter EOS from partial decode)
            text_so_far = self._tokenizer.decode(
                [t for t in generated if t != self.config.eos_token_id]
            )
            if text_so_far != prev_text:
                yield text_so_far[len(prev_text) :]
                prev_text = text_so_far

            if token == self.config.eos_token_id or len(generated) > max_tokens:
                break

            # Encode more audio chunks on demand
            if enc_chunk_gen is not None and pos >= adapter_len:
                try:
                    chunk_out = next(enc_chunk_gen)
                    chunk_adapter = self.encoder.downsample_and_project(chunk_out)
                    mx.eval(chunk_adapter)
                    adapter_out = mx.concatenate([adapter_out, chunk_adapter], axis=0)
                    adapter_len = adapter_out.shape[0]
                except StopIteration:
                    enc_chunk_gen = None

            if pos < adapter_len:
                embed = adapter_out[pos] + self.decoder.embed_token(token)
            else:
                embed = self.decoder.embed_token(token)

            h, cache = self.decoder.forward(embed[None, :], start_pos=pos, cache=cache)
            logits = self.decoder.logits(h.squeeze(0))
            next_tok = self._next_token_mx(logits, temperature)
            mx.async_eval(next_tok)

            if len(generated) % 256 == 0:
                mx.clear_cache()
        else:
            # Loop completed without break — read final pending token
            token = int(next_tok.item())
            generated.append(token)
            text_so_far = self._tokenizer.decode(
                [t for t in generated if t != self.config.eos_token_id]
            )
            if text_so_far != prev_text:
                yield text_so_far[len(prev_text) :]

        mx.clear_cache()

    def _next_token_mx(self, logits, temperature):
        """Compute next token as lazy mx.array (for async eval pipelining)."""
        if temperature == 0:
            return mx.argmax(logits)
        return mx.random.categorical(logits * (1.0 / temperature))

    def _sample(self, logits, temperature):
        # categorical accepts unnormalized logits directly — no softmax+log needed
        return int(self._next_token_mx(logits, temperature).item())

    def sanitize(self, weights):
        """Map weight names from consolidated.safetensors to our module structure."""
        new_weights = {}

        enc_prefix = "mm_streams_embeddings.embedding_module.whisper_encoder"
        adapter_prefix = "mm_streams_embeddings.embedding_module"
        tok_emb_key = "mm_streams_embeddings.embedding_module.tok_embeddings.weight"

        for k, v in weights.items():
            new_key = None

            if k == tok_emb_key:
                new_key = "decoder.tok_embeddings.weight"

            elif k == "norm.weight":
                new_key = "decoder.norm.weight"

            elif k.startswith(f"{enc_prefix}.conv_layers."):
                # e.g., ...conv_layers.0.conv.weight -> encoder.conv_layers_0_conv.conv.weight
                rest = k[len(f"{enc_prefix}.conv_layers.") :]
                # rest: "0.conv.weight" or "0.conv.bias"
                parts = rest.split(".", 2)  # ['0', 'conv', 'weight']
                layer_idx = parts[0]
                param = parts[2]  # 'weight' or 'bias'
                new_key = f"encoder.conv_layers_{layer_idx}_conv.conv.{param}"

                # Transpose conv weights from PyTorch [out, in, k] to MLX [out, k, in]
                if param == "weight" and v.ndim == 3:
                    v = v.transpose(0, 2, 1)

            elif k.startswith(f"{enc_prefix}.transformer.layers."):
                rest = k[len(f"{enc_prefix}.transformer.layers.") :]
                # e.g., "0.attention.wq.weight"
                parts = rest.split(".", 1)
                layer_idx = parts[0]
                param_path = parts[1]

                # Map FFN weights
                param_path = param_path.replace("feed_forward.w1.", "feed_forward_w1.")
                param_path = param_path.replace("feed_forward.w2.", "feed_forward_w2.")
                param_path = param_path.replace("feed_forward.w3.", "feed_forward_w3.")

                new_key = f"encoder.transformer_layers.{layer_idx}.{param_path}"

            elif k.startswith(f"{enc_prefix}.transformer.norm."):
                rest = k[len(f"{enc_prefix}.transformer.norm.") :]
                new_key = f"encoder.transformer_norm.{rest}"

            elif k.startswith(f"{adapter_prefix}.audio_language_projection."):
                rest = k[len(f"{adapter_prefix}.audio_language_projection.") :]
                # "0.weight" -> "audio_language_projection_0.weight"
                # "2.weight" -> "audio_language_projection_2.weight"
                parts = rest.split(".", 1)
                idx = parts[0]
                param = parts[1]
                new_key = f"encoder.audio_language_projection_{idx}.{param}"

            elif k.startswith("layers."):
                rest = k[len("layers.") :]
                # e.g., "0.attention.wq.weight"
                parts = rest.split(".", 1)
                layer_idx = parts[0]
                param_path = parts[1]

                # Map FFN
                param_path = param_path.replace("feed_forward.w1.", "feed_forward_w1.")
                param_path = param_path.replace("feed_forward.w2.", "feed_forward_w2.")
                param_path = param_path.replace("feed_forward.w3.", "feed_forward_w3.")
                # Map ada norm: ada_rms_norm_t_cond.0.weight -> ada_rms_norm_t_cond.ada_down.weight
                param_path = param_path.replace(
                    "ada_rms_norm_t_cond.0.", "ada_rms_norm_t_cond.ada_down."
                )
                param_path = param_path.replace(
                    "ada_rms_norm_t_cond.2.", "ada_rms_norm_t_cond.ada_up."
                )

                new_key = f"decoder.layers.{layer_idx}.{param_path}"

            if new_key is not None:
                new_weights[new_key] = v
            else:
                # Pass through any unrecognized weights as-is
                new_weights[k] = v

        return new_weights

    def model_quant_predicate(self, p, m):
        """Quantize all big linears; only skip the small ones that hurt quality.

        Per-step memory bandwidth dominates decode time, so the tied
        ``tok_embeddings`` matmul in ``decoder.logits`` and the audio adapter
        projections MUST be quantized to match voxmlx's throughput. Only the
        norms, the tiny AdaRMSNorm MLPs, and the conv stem (3×128 kernels) are
        skipped — matmul kernels on those are negligible and quantizing them
        would cost more in quality than it saves in bandwidth.
        """
        skip_patterns = ["norm", "ada_rms_norm", "conv_layers"]
        return not any(pat in p for pat in skip_patterns)

    @classmethod
    def post_load_hook(cls, model: "Model", model_path: Path) -> "Model":
        """Initialize tokenizer and precompute ada scales after weight loading."""
        model_path = Path(model_path)

        # Load Tekken tokenizer
        model._tokenizer = TekkenTokenizer.from_model_path(model_path)

        # Precompute mel filters
        model._ensure_mel_filters()

        # No dtype/quantization coercion here: the framework's
        # ``apply_quantization`` already decides per-module whether to quantize
        # based on whether ``<path>.scales`` is present in the checkpoint
        # (see mlx_audio.utils.apply_quantization). That makes the model adapt
        # automatically — an old fp16 / unquantized-embedding checkpoint keeps
        # its old layout (and old perf), a new bf16 / 4-bit-embedding
        # checkpoint runs at full speed. Any dtype upgrade belongs in the
        # converter, not at load time.
        n_delay = _num_delay_tokens(model.config.transcription_delay_ms)
        t_cond = compute_time_embedding(float(n_delay), model.config.decoder.dim)
        # Match whatever dtype the AdaRMSNorm weights ended up in.
        ada_weight = model.decoder.layers[0].ada_rms_norm_t_cond.ada_down.weight
        t_cond = t_cond.astype(ada_weight.dtype)
        model.decoder.precompute_ada_scales(t_cond)
        model._ada_scale_delay = n_delay
        # Evaluate ada scales eagerly
        for scale in model.decoder._ada_scales:
            if scale is not None:
                mx.eval(scale)

        return model