//! Moonshine v2 candle port — encoder-decoder ASR transformer designed
//! for low-latency English transcription.
//!
//! See `docs/moonshine_port_notes.md` for the architectural plan, the
//! tensor layout from `examples/moonshine_inspect`, and the porting
//! task list.
//!
//! ## Status
//!
//! - [x] Config + Conv1d audio stem (this commit, Phase 8.5)
//! - [ ] Encoder transformer layer (partial RoPE, GELU MLP)
//! - [ ] Decoder transformer layer (self-attn + cross-attn, SwiGLU MLP)
//! - [ ] Generation loop (greedy + KV cache)
//! - [ ] Tokenizer wiring (HF `tokenizer.json` via `tokenizers` crate)
//! - [ ] Weight mapping + smoke test against HF reference output
//!
//! Each item is roughly an hour of focused work; total ~12-15 h.
//!
//! ## Architecture (from HF `UsefulSensors/moonshine-tiny/config.json`)
//!
//! - hidden_size = 288, intermediate_size = 1152
//! - 6 encoder layers + 6 decoder layers, 8 heads each (head_dim = 36,
//!   pad to 8 → 40 — see `pad_head_dim_to_multiple_of`)
//! - vocab_size = 32_768, max_position_embeddings = 194 (decoder text)
//! - encoder MLP: GELU, no gating
//! - decoder MLP: SiLU SwiGLU (fc1 outputs 2× intermediate, split gate/up)
//! - RoPE: theta=10000, partial_rotary_factor=0.9 (rotary on first 90 %
//!   of head_dim)
//! - bos = 1, eos = 2, pad = 2, decoder_start = 1

use candle_core::{DType, Device, Module, Result, Tensor};
use candle_nn::{Conv1d, Conv1dConfig, LayerNorm, Linear, VarBuilder};

/// Hyperparameters matching `UsefulSensors/moonshine-tiny/config.json`.
#[derive(Debug, Clone)]
pub struct MoonshineConfig {
    pub hidden_size: usize,
    pub intermediate_size: usize,
    pub vocab_size: usize,
    pub max_position_embeddings: usize,
    pub encoder_num_hidden_layers: usize,
    pub encoder_num_attention_heads: usize,
    pub decoder_num_hidden_layers: usize,
    pub decoder_num_attention_heads: usize,
    pub rope_theta: f64,
    pub partial_rotary_factor: f64,
    pub pad_head_dim_to_multiple_of: usize,
    pub bos_token_id: u32,
    pub eos_token_id: u32,
    pub decoder_start_token_id: u32,
}

impl MoonshineConfig {
    /// Tiny variant (27.1 M params).
    pub fn tiny() -> Self {
        Self {
            hidden_size: 288,
            intermediate_size: 1152,
            vocab_size: 32_768,
            max_position_embeddings: 194,
            encoder_num_hidden_layers: 6,
            encoder_num_attention_heads: 8,
            decoder_num_hidden_layers: 6,
            decoder_num_attention_heads: 8,
            rope_theta: 10000.0,
            partial_rotary_factor: 0.9,
            pad_head_dim_to_multiple_of: 8,
            bos_token_id: 1,
            eos_token_id: 2,
            decoder_start_token_id: 1,
        }
    }

    /// Padded head dim (the actual size used for RoPE + attention math).
    /// `head_dim_raw = hidden_size / num_heads`, rounded up to the
    /// nearest multiple of `pad_head_dim_to_multiple_of`.
    pub fn padded_head_dim(&self, num_heads: usize) -> usize {
        let raw = self.hidden_size / num_heads;
        let pad = self.pad_head_dim_to_multiple_of;
        raw.div_ceil(pad) * pad
    }
}

/// Three-layer Conv1d audio stem mapping raw 16 kHz waveform to a
/// sequence of 288-d hidden vectors at ~42 Hz frame rate.
///
/// Layout (from HF `modeling_moonshine.py`):
/// - conv1: in=1, out=288, kernel=127, stride=64, no bias
/// - conv2: in=288, out=576, kernel=7, stride=3, bias
/// - conv3: in=576, out=288, kernel=3, stride=2, bias
///
/// Total downsampling: 64 × 3 × 2 = 384 ×.
/// 16 kHz audio → output ~42 Hz frame rate (≈ 24 ms / frame).
///
/// Activations between convs: HF source uses GELU after conv1 and a
/// LayerNorm + GELU after conv2 (TODO: verify the exact ordering when
/// the encoder transformer block is wired).
pub struct ConvStem {
    conv1: Conv1d,
    conv2: Conv1d,
    conv3: Conv1d,
}

impl ConvStem {
    pub fn new(vb: VarBuilder) -> Result<Self> {
        let conv1 = candle_nn::conv1d_no_bias(
            1,
            288,
            127,
            Conv1dConfig {
                padding: 0,
                stride: 64,
                dilation: 1,
                groups: 1,
                cudnn_fwd_algo: None,
            },
            vb.pp("conv1"),
        )?;
        let conv2 = candle_nn::conv1d(
            288,
            576,
            7,
            Conv1dConfig {
                padding: 0,
                stride: 3,
                dilation: 1,
                groups: 1,
                cudnn_fwd_algo: None,
            },
            vb.pp("conv2"),
        )?;
        let conv3 = candle_nn::conv1d(
            576,
            288,
            3,
            Conv1dConfig {
                padding: 0,
                stride: 2,
                dilation: 1,
                groups: 1,
                cudnn_fwd_algo: None,
            },
            vb.pp("conv3"),
        )?;
        Ok(Self {
            conv1,
            conv2,
            conv3,
        })
    }

    /// Forward: `(B, 1, T_audio)` raw waveform → `(B, T_seq, 288)` where
    /// `T_seq ≈ T_audio / 384`. Returns hidden states ready for the
    /// encoder transformer stack.
    pub fn forward(&self, pcm: &Tensor) -> Result<Tensor> {
        // Conv1: raw audio → 288 channels at ~250 Hz
        let h = self.conv1.forward(pcm)?;
        let h = h.tanh()?; // HF source uses tanh after conv1
        // Conv2: 288 → 576, stride 3
        let h = self.conv2.forward(&h)?;
        let h = h.gelu_erf()?;
        // Conv3: 576 → 288, stride 2
        let h = self.conv3.forward(&h)?;
        let h = h.gelu_erf()?;
        // Output is (B, 288, T_seq); transpose to (B, T_seq, 288) for
        // the transformer encoder layers.
        h.transpose(1, 2)?.contiguous()
    }
}

/// Loader: open the HF safetensors and construct a `ConvStem`. Useful
/// for the standalone Phase 8.5 smoke test.
pub fn load_conv_stem(weights_path: &std::path::Path, device: &Device) -> Result<ConvStem> {
    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, device) }?;
    ConvStem::new(vb.pp("model").pp("encoder"))
}

// ---------------------------------------------------------------------
// Phase 8.6 — encoder transformer block
// ---------------------------------------------------------------------

/// Cosine/sine tables for partial RoPE.
///
/// Moonshine uses `partial_rotary_factor = 0.9` with `head_dim = 36`
/// (288 / 8 heads). `rotary_dim = floor(36 * 0.9 / 2) * 2 = 32` (must
/// be even — RoPE rotates pairs). The remaining 4 dims pass through
/// unchanged. RoPE base is `theta = 10_000`.
struct RotaryCache {
    cos: Tensor, // (max_seq, rotary_dim/2)
    sin: Tensor, // (max_seq, rotary_dim/2)
    rotary_dim: usize,
}

impl RotaryCache {
    fn new(
        rotary_dim: usize,
        max_seq: usize,
        theta: f64,
        dtype: DType,
        dev: &Device,
    ) -> Result<Self> {
        assert!(rotary_dim.is_multiple_of(2), "rotary_dim must be even");
        let inv_freq: Vec<f32> = (0..rotary_dim)
            .step_by(2)
            .map(|i| (1.0 / theta.powf(i as f64 / rotary_dim as f64)) as f32)
            .collect();
        let inv_freq = Tensor::new(inv_freq, dev)?;
        let positions = Tensor::arange(0u32, max_seq as u32, dev)?
            .to_dtype(DType::F32)?
            .reshape((max_seq, 1))?;
        let freqs = positions.matmul(&inv_freq.reshape((1, rotary_dim / 2))?)?;
        let cos = freqs.cos()?.to_dtype(dtype)?;
        let sin = freqs.sin()?.to_dtype(dtype)?;
        Ok(Self {
            cos,
            sin,
            rotary_dim,
        })
    }

    /// Apply partial RoPE to `q` of shape `(B, H, T, head_dim)`.
    /// Rotates the first `rotary_dim` channels; leaves the rest as-is.
    /// Uses positions `0..T` (prefill mode).
    fn apply(&self, x: &Tensor) -> Result<Tensor> {
        let (_b, _h, t, _) = x.dims4()?;
        self.apply_at(x, 0, t)
    }

    /// Apply RoPE for a window starting at `position`, length `t`.
    /// Used by the cached single-token step (`t=1`, `position=cache_len`).
    fn apply_at(&self, x: &Tensor, position: usize, t: usize) -> Result<Tensor> {
        let (_b, _h, _t, head_dim) = x.dims4()?;
        let cos = self.cos.narrow(0, position, t)?;
        let sin = self.sin.narrow(0, position, t)?;
        if head_dim == self.rotary_dim {
            return candle_nn::rotary_emb::rope_i(x, &cos, &sin);
        }
        let rot = x.narrow(3, 0, self.rotary_dim)?.contiguous()?;
        let pass = x.narrow(3, self.rotary_dim, head_dim - self.rotary_dim)?;
        let rot = candle_nn::rotary_emb::rope_i(&rot, &cos, &sin)?;
        Tensor::cat(&[&rot, &pass], 3)
    }
}

/// Self-attention block matching `model.encoder.layers.X.self_attn`.
/// 8 heads × 36 head_dim, no bias (`attention_bias=false`), partial
/// RoPE on first 32 head_dim channels.
struct EncoderAttention {
    q_proj: Linear,
    k_proj: Linear,
    v_proj: Linear,
    o_proj: Linear,
    n_heads: usize,
    head_dim: usize,
}

impl EncoderAttention {
    fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
        let h = cfg.hidden_size;
        let n_heads = cfg.encoder_num_attention_heads;
        let head_dim = h / n_heads;
        let q_proj = candle_nn::linear_no_bias(h, h, vb.pp("q_proj"))?;
        let k_proj = candle_nn::linear_no_bias(h, h, vb.pp("k_proj"))?;
        let v_proj = candle_nn::linear_no_bias(h, h, vb.pp("v_proj"))?;
        let o_proj = candle_nn::linear_no_bias(h, h, vb.pp("o_proj"))?;
        Ok(Self {
            q_proj,
            k_proj,
            v_proj,
            o_proj,
            n_heads,
            head_dim,
        })
    }

    fn forward(&self, xs: &Tensor, rope: &RotaryCache) -> Result<Tensor> {
        let (b, t, _h) = xs.dims3()?;
        let q = self.q_proj.forward(xs)?;
        let k = self.k_proj.forward(xs)?;
        let v = self.v_proj.forward(xs)?;
        // (B, T, H) -> (B, H_heads, T, head_dim)
        let q = q
            .reshape((b, t, self.n_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?;
        let k = k
            .reshape((b, t, self.n_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?;
        let v = v
            .reshape((b, t, self.n_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?;
        let q = rope.apply(&q)?;
        let k = rope.apply(&k)?;
        // Collapse (B, H, T, D) -> (B*H, T, D) for the matmul. candle's
        // 4D batched matmul on Metal had shape-mismatch issues for our
        // pattern; the 3D form is unambiguous.
        let bh = b * self.n_heads;
        let q3 = q.reshape((bh, t, self.head_dim))?;
        let k3 = k.reshape((bh, t, self.head_dim))?;
        let v3 = v.reshape((bh, t, self.head_dim))?;
        let scale = 1.0 / (self.head_dim as f64).sqrt();
        let scores = (q3.matmul(&k3.transpose(1, 2)?.contiguous()?)? * scale)?;
        let probs = candle_nn::ops::softmax_last_dim(&scores)?;
        let out = probs.matmul(&v3)?;
        // (B*H, T, D) -> (B, H, T, D) -> (B, T, H*D)
        let out = out
            .reshape((b, self.n_heads, t, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?
            .reshape((b, t, self.n_heads * self.head_dim))?;
        self.o_proj.forward(&out)
    }
}

/// FFN block matching `model.encoder.layers.X.mlp`.
/// 288 → 1152 (with bias) → GELU → 1152 → 288 (with bias).
struct EncoderMlp {
    fc1: Linear,
    fc2: Linear,
}

impl EncoderMlp {
    fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
        let fc1 = candle_nn::linear(cfg.hidden_size, cfg.intermediate_size, vb.pp("fc1"))?;
        let fc2 = candle_nn::linear(cfg.intermediate_size, cfg.hidden_size, vb.pp("fc2"))?;
        Ok(Self { fc1, fc2 })
    }

    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
        let h = self.fc1.forward(xs)?;
        let h = h.gelu_erf()?;
        self.fc2.forward(&h)
    }
}

/// One encoder layer: Pre-LN attn + Pre-LN MLP, residuals as standard.
/// Tensors are `LayerNorm` weight-only (no bias) per the inspector
/// dump — implemented via `LayerNorm` with bias zeros.
struct EncoderLayer {
    input_ln: LayerNorm,
    self_attn: EncoderAttention,
    post_attn_ln: LayerNorm,
    mlp: EncoderMlp,
}

impl EncoderLayer {
    fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
        let h = cfg.hidden_size;
        let input_ln = layer_norm_weight_only(h, 1e-5, vb.pp("input_layernorm"))?;
        let self_attn = EncoderAttention::new(cfg, vb.pp("self_attn"))?;
        let post_attn_ln = layer_norm_weight_only(h, 1e-5, vb.pp("post_attention_layernorm"))?;
        let mlp = EncoderMlp::new(cfg, vb.pp("mlp"))?;
        Ok(Self {
            input_ln,
            self_attn,
            post_attn_ln,
            mlp,
        })
    }

    fn forward(&self, xs: &Tensor, rope: &RotaryCache) -> Result<Tensor> {
        let h = self.input_ln.forward(xs)?;
        let h = self.self_attn.forward(&h, rope)?;
        let xs = (xs + h)?;
        let h = self.post_attn_ln.forward(&xs)?;
        let h = self.mlp.forward(&h)?;
        xs + h
    }
}

/// Helper: `LayerNorm` with weight only (zeros bias). Moonshine's
/// `*_layernorm.weight [288]` tensors lack a corresponding `.bias`.
fn layer_norm_weight_only(size: usize, eps: f64, vb: VarBuilder) -> Result<LayerNorm> {
    let weight = vb.get((size,), "weight")?;
    let bias = Tensor::zeros((size,), weight.dtype(), weight.device())?;
    Ok(LayerNorm::new(weight, bias, eps))
}

/// Full encoder: conv stem → 6 transformer layers → final layer norm.
/// Output shape: `(B, T_seq, hidden=288)` ready to feed cross-attn in
/// the decoder (Phase 8.7).
pub struct Encoder {
    stem: ConvStem,
    layers: Vec<EncoderLayer>,
    final_ln: LayerNorm,
    rope: RotaryCache,
}

impl Encoder {
    pub fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
        let stem = ConvStem::new(vb.pp("encoder"))?;
        let head_dim = cfg.hidden_size / cfg.encoder_num_attention_heads;
        let rotary_dim = ((head_dim as f64 * cfg.partial_rotary_factor) as usize / 2) * 2;
        // max_seq for the encoder is bounded by the input audio length /
        // total stride 384. For 30 s of 16 kHz audio that's ~1250.
        // Allocate a generous table.
        let rope = RotaryCache::new(rotary_dim, 4096, cfg.rope_theta, vb.dtype(), vb.device())?;
        let mut layers = Vec::with_capacity(cfg.encoder_num_hidden_layers);
        let layer_vb = vb.pp("encoder").pp("layers");
        for i in 0..cfg.encoder_num_hidden_layers {
            layers.push(EncoderLayer::new(cfg, layer_vb.pp(i))?);
        }
        let final_ln =
            layer_norm_weight_only(cfg.hidden_size, 1e-5, vb.pp("encoder").pp("layer_norm"))?;
        Ok(Self {
            stem,
            layers,
            final_ln,
            rope,
        })
    }

    pub fn forward(&self, pcm: &Tensor) -> Result<Tensor> {
        let mut h = self.stem.forward(pcm)?;
        for layer in &self.layers {
            h = layer.forward(&h, &self.rope)?;
        }
        self.final_ln.forward(&h)
    }
}

/// Loader: open the HF safetensors and construct an `Encoder` with
/// weights at `model.*` (the canonical HF prefix).
pub fn load_encoder(
    weights_path: &std::path::Path,
    device: &Device,
    cfg: &MoonshineConfig,
) -> Result<Encoder> {
    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, device) }?;
    Encoder::new(cfg, vb.pp("model"))
}

// ---------------------------------------------------------------------
// Phase 8.7 — decoder transformer block
// ---------------------------------------------------------------------

/// Cross-attention block matching `model.decoder.layers.X.encoder_attn`.
/// Queries come from the decoder hidden state; keys and values come
/// from the encoder output (computed once per generation step, cached
/// across generation tokens — caching deferred to Phase 8.8).
struct CrossAttention {
    q_proj: Linear,
    k_proj: Linear,
    v_proj: Linear,
    o_proj: Linear,
    n_heads: usize,
    head_dim: usize,
}

impl CrossAttention {
    fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
        let h = cfg.hidden_size;
        let n_heads = cfg.decoder_num_attention_heads;
        let head_dim = h / n_heads;
        Ok(Self {
            q_proj: candle_nn::linear_no_bias(h, h, vb.pp("q_proj"))?,
            k_proj: candle_nn::linear_no_bias(h, h, vb.pp("k_proj"))?,
            v_proj: candle_nn::linear_no_bias(h, h, vb.pp("v_proj"))?,
            o_proj: candle_nn::linear_no_bias(h, h, vb.pp("o_proj"))?,
            n_heads,
            head_dim,
        })
    }

    /// `xs`: decoder hidden `(B, T_dec, H)`.
    /// `enc`: encoder output `(B, T_enc, H)`.
    /// Returns: `(B, T_dec, H)`.
    fn forward(&self, xs: &Tensor, enc: &Tensor) -> Result<Tensor> {
        let (b, t_dec, _h) = xs.dims3()?;
        let (_, t_enc, _) = enc.dims3()?;
        let q = self.q_proj.forward(xs)?;
        let k = self.k_proj.forward(enc)?;
        let v = self.v_proj.forward(enc)?;
        // (B, T, H) -> (B*H_heads, T, head_dim) — same 3D matmul
        // pattern as EncoderAttention to dodge the candle 4D Metal bug.
        let bh = b * self.n_heads;
        let q = q
            .reshape((b, t_dec, self.n_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?
            .reshape((bh, t_dec, self.head_dim))?;
        let k = k
            .reshape((b, t_enc, self.n_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?
            .reshape((bh, t_enc, self.head_dim))?;
        let v = v
            .reshape((b, t_enc, self.n_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?
            .reshape((bh, t_enc, self.head_dim))?;
        let scale = 1.0 / (self.head_dim as f64).sqrt();
        // Cross-attention has NO causal mask: each decoder token can
        // attend to every encoder position.
        let scores = (q.matmul(&k.transpose(1, 2)?.contiguous()?)? * scale)?;
        let probs = candle_nn::ops::softmax_last_dim(&scores)?;
        let out = probs.matmul(&v)?;
        let out = out
            .reshape((b, self.n_heads, t_dec, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?
            .reshape((b, t_dec, self.n_heads * self.head_dim))?;
        self.o_proj.forward(&out)
    }
}

/// Decoder self-attention. Same shape as encoder attention but applies
/// a causal mask so token `i` only attends to tokens `0..=i`. Partial
/// RoPE on q/k as in the encoder.
struct DecoderSelfAttention {
    q_proj: Linear,
    k_proj: Linear,
    v_proj: Linear,
    o_proj: Linear,
    n_heads: usize,
    head_dim: usize,
}

impl DecoderSelfAttention {
    fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
        let h = cfg.hidden_size;
        let n_heads = cfg.decoder_num_attention_heads;
        let head_dim = h / n_heads;
        Ok(Self {
            q_proj: candle_nn::linear_no_bias(h, h, vb.pp("q_proj"))?,
            k_proj: candle_nn::linear_no_bias(h, h, vb.pp("k_proj"))?,
            v_proj: candle_nn::linear_no_bias(h, h, vb.pp("v_proj"))?,
            o_proj: candle_nn::linear_no_bias(h, h, vb.pp("o_proj"))?,
            n_heads,
            head_dim,
        })
    }

    fn forward(&self, xs: &Tensor, rope: &RotaryCache) -> Result<Tensor> {
        let (b, t, _h) = xs.dims3()?;
        let q = self.q_proj.forward(xs)?;
        let k = self.k_proj.forward(xs)?;
        let v = self.v_proj.forward(xs)?;
        let q = q
            .reshape((b, t, self.n_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?;
        let k = k
            .reshape((b, t, self.n_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?;
        let v = v
            .reshape((b, t, self.n_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?;
        let q = rope.apply(&q)?;
        let k = rope.apply(&k)?;
        let bh = b * self.n_heads;
        let q3 = q.reshape((bh, t, self.head_dim))?;
        let k3 = k.reshape((bh, t, self.head_dim))?;
        let v3 = v.reshape((bh, t, self.head_dim))?;
        let scale = 1.0 / (self.head_dim as f64).sqrt();
        let mut scores = (q3.matmul(&k3.transpose(1, 2)?.contiguous()?)? * scale)?;
        // Causal mask: build a (T, T) upper-triangular -inf mask.
        let mask = causal_mask(t, scores.dtype(), scores.device())?;
        scores = scores.broadcast_add(&mask)?;
        let probs = candle_nn::ops::softmax_last_dim(&scores)?;
        let out = probs.matmul(&v3)?;
        let out = out
            .reshape((b, self.n_heads, t, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?
            .reshape((b, t, self.n_heads * self.head_dim))?;
        self.o_proj.forward(&out)
    }
}

fn causal_mask(t: usize, dtype: DType, device: &Device) -> Result<Tensor> {
    // Upper-triangular `-inf` mask: value at (i, j) is 0 if j <= i, else -inf.
    let mut data = vec![0.0f32; t * t];
    for i in 0..t {
        for j in (i + 1)..t {
            data[i * t + j] = f32::NEG_INFINITY;
        }
    }
    Tensor::from_vec(data, (t, t), device)?.to_dtype(dtype)
}

/// SwiGLU MLP for the decoder. `fc1` is the fused gate+up projection:
/// weight shape `[2 * intermediate_size, hidden_size] = [2304, 288]`.
/// Output is split on dim -1 into `gate` and `up`, both `[..., 1152]`.
/// Activation: `silu(gate) * up`. Then `fc2: [hidden, intermediate] =
/// [288, 1152]` projects back to `hidden`.
struct DecoderMlp {
    fc1: Linear,
    fc2: Linear,
    intermediate: usize,
}

impl DecoderMlp {
    fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
        // fc1 has bias; fc2 has bias (consistent with encoder MLP — the
        // inspector dump shows `.bias` on both).
        let fc1 = candle_nn::linear(cfg.hidden_size, cfg.intermediate_size * 2, vb.pp("fc1"))?;
        let fc2 = candle_nn::linear(cfg.intermediate_size, cfg.hidden_size, vb.pp("fc2"))?;
        Ok(Self {
            fc1,
            fc2,
            intermediate: cfg.intermediate_size,
        })
    }

    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
        // (B, T, H) -> (B, T, 2 * intermediate)
        let h = self.fc1.forward(xs)?;
        let dims = h.dims();
        let last = dims.len() - 1;
        // HF Moonshine MLP (verified against modeling_moonshine.py):
        //   hidden, gate = fc1(x).chunk(2, dim=-1)
        //   out = silu(gate) * hidden
        // i.e., FIRST half is `up` (hidden), SECOND half is `gate`.
        // We had this reversed initially, which produced a degenerate
        // repetition loop after 1-2 tokens.
        let up = h.narrow(last, 0, self.intermediate)?;
        let gate = h.narrow(last, self.intermediate, self.intermediate)?;
        let activated = candle_nn::ops::silu(&gate)?.mul(&up)?;
        self.fc2.forward(&activated)
    }
}

/// One decoder layer: Pre-LN self-attn, Pre-LN cross-attn, Pre-LN MLP.
/// Three layer norms per layer (input_layernorm, post_attention_layernorm,
/// final_layernorm), all weight-only.
struct DecoderLayer {
    input_ln: LayerNorm,
    self_attn: DecoderSelfAttention,
    post_attn_ln: LayerNorm,
    cross_attn: CrossAttention,
    final_ln: LayerNorm,
    mlp: DecoderMlp,
}

impl DecoderLayer {
    fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
        let h = cfg.hidden_size;
        Ok(Self {
            input_ln: layer_norm_weight_only(h, 1e-5, vb.pp("input_layernorm"))?,
            self_attn: DecoderSelfAttention::new(cfg, vb.pp("self_attn"))?,
            post_attn_ln: layer_norm_weight_only(h, 1e-5, vb.pp("post_attention_layernorm"))?,
            cross_attn: CrossAttention::new(cfg, vb.pp("encoder_attn"))?,
            final_ln: layer_norm_weight_only(h, 1e-5, vb.pp("final_layernorm"))?,
            mlp: DecoderMlp::new(cfg, vb.pp("mlp"))?,
        })
    }

    fn forward(&self, xs: &Tensor, enc: &Tensor, rope: &RotaryCache) -> Result<Tensor> {
        // Self-attn (causal)
        let h = self.input_ln.forward(xs)?;
        let h = self.self_attn.forward(&h, rope)?;
        let xs = (xs + h)?;
        // Cross-attn
        let h = self.post_attn_ln.forward(&xs)?;
        let h = self.cross_attn.forward(&h, enc)?;
        let xs = (xs + h)?;
        // MLP
        let h = self.final_ln.forward(&xs)?;
        let h = self.mlp.forward(&h)?;
        xs + h
    }
}

/// Full decoder: token embedding → 6 layers → final LN → tied LM head.
pub struct Decoder {
    embed: candle_nn::Embedding,
    layers: Vec<DecoderLayer>,
    final_ln: LayerNorm,
    rope: RotaryCache,
    /// Cached transposed embedding for the tied LM head.
    /// Computing logits: hidden @ embed.weight.T -> (B, T, vocab).
    embed_weight: Tensor,
}

impl Decoder {
    pub fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
        let embed = candle_nn::embedding(
            cfg.vocab_size,
            cfg.hidden_size,
            vb.pp("decoder").pp("embed_tokens"),
        )?;
        let head_dim = cfg.hidden_size / cfg.decoder_num_attention_heads;
        let rotary_dim = ((head_dim as f64 * cfg.partial_rotary_factor) as usize / 2) * 2;
        let rope = RotaryCache::new(
            rotary_dim,
            cfg.max_position_embeddings.max(512),
            cfg.rope_theta,
            vb.dtype(),
            vb.device(),
        )?;
        let mut layers = Vec::with_capacity(cfg.decoder_num_hidden_layers);
        let layer_vb = vb.pp("decoder").pp("layers");
        for i in 0..cfg.decoder_num_hidden_layers {
            layers.push(DecoderLayer::new(cfg, layer_vb.pp(i))?);
        }
        // Decoder uses `decoder.norm.weight` (the encoder uses
        // `encoder.layer_norm.weight`). Inspector dump confirmed.
        let final_ln = layer_norm_weight_only(cfg.hidden_size, 1e-5, vb.pp("decoder").pp("norm"))?;
        let embed_weight = embed.embeddings().clone();
        Ok(Self {
            embed,
            layers,
            final_ln,
            rope,
            embed_weight,
        })
    }

    /// Forward over `tokens` shape `(B, T)` with encoder output `enc`
    /// shape `(B, T_enc, H)`. Returns logits `(B, T, vocab_size)`.
    /// **No KV cache** — this is the prefill / single-step path.
    /// The full streaming-generation loop with KV cache lands in
    /// Phase 8.8.
    pub fn forward(&self, tokens: &Tensor, enc: &Tensor) -> Result<Tensor> {
        let mut h = self.embed.forward(tokens)?;
        for layer in &self.layers {
            h = layer.forward(&h, enc, &self.rope)?;
        }
        h = self.final_ln.forward(&h)?;
        // LM head tied to embedding: logits = hidden @ embed.T.
        // embed.weight shape: (vocab, hidden). transpose -> (hidden, vocab).
        let lm_w = self.embed_weight.transpose(0, 1)?.contiguous()?;
        h.broadcast_matmul(&lm_w)
    }
}

/// Loader: open the HF safetensors and construct an encoder + decoder.
pub fn load_full(
    weights_path: &std::path::Path,
    device: &Device,
    cfg: &MoonshineConfig,
) -> Result<(Encoder, Decoder)> {
    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, device) }?;
    let encoder = Encoder::new(cfg, vb.pp("model"))?;
    let decoder = Decoder::new(cfg, vb.pp("model"))?;
    Ok((encoder, decoder))
}

// ---------------------------------------------------------------------
// Phase 8.8 — generation loop + tokenizer
// ---------------------------------------------------------------------

impl Decoder {
    /// Greedy autoregressive decode given an encoder output. Stops on
    /// `eos_token_id` or after `max_tokens`. Returns the predicted
    /// token ids (excluding the initial start token).
    ///
    /// **No KV cache** — each step re-runs the decoder on the full
    /// growing token sequence (O(T^2) total work). For Moonshine-tiny
    /// at ~85 ms / step prefill, a 30-token transcript takes a few
    /// seconds. KV-cached `step()` is a follow-up if needed.
    pub fn generate(
        &self,
        encoder_output: &Tensor,
        cfg: &MoonshineConfig,
        max_tokens: usize,
    ) -> Result<Vec<u32>> {
        let device = encoder_output.device();
        let mut tokens: Vec<u32> = vec![cfg.decoder_start_token_id];
        let mut out = Vec::with_capacity(max_tokens);
        for _ in 0..max_tokens {
            let input = Tensor::from_vec(tokens.clone(), (1, tokens.len()), device)?;
            let logits = self.forward(&input, encoder_output)?;
            // logits: (1, T_so_far, vocab). Take last position.
            let last_t = tokens.len() - 1;
            let last = logits.narrow(1, last_t, 1)?.squeeze(1)?; // (1, vocab)
            let argmax = last.argmax(1)?;
            let next_id: u32 = argmax.to_dtype(DType::U32)?.to_vec1::<u32>()?[0];
            if next_id == cfg.eos_token_id {
                break;
            }
            out.push(next_id);
            tokens.push(next_id);
            if tokens.len() >= cfg.max_position_embeddings {
                break;
            }
        }
        Ok(out)
    }
}

/// Convenience: load the HF tokenizer.json for Moonshine. Caller passes
/// the path returned by hf_hub.
pub fn load_tokenizer(
    path: &std::path::Path,
) -> std::result::Result<tokenizers::Tokenizer, Box<dyn std::error::Error + Send + Sync>> {
    tokenizers::Tokenizer::from_file(path)
}

// ---------------------------------------------------------------------
// Phase 8.9 — KV cache for single-token autoregressive decoding
// ---------------------------------------------------------------------

/// Per-layer K/V cache for one decode session.
/// - `self_k[i]` / `self_v[i]`: cumulative self-attn K/V for layer i,
///   shape `(B*H, position, head_dim)`. Grows by 1 per `step()`.
/// - `cross_k[i]` / `cross_v[i]`: encoder cross-attn K/V for layer i,
///   computed once via `precompute_cross_kv` and reused every step.
pub struct DecoderCache {
    self_k: Vec<Option<Tensor>>,
    self_v: Vec<Option<Tensor>>,
    cross_k: Vec<Tensor>,
    cross_v: Vec<Tensor>,
    pub position: usize,
}

impl DecoderSelfAttention {
    /// Single-token cached self-attention. Appends new K/V to the
    /// cache and attends across the full accumulated history.
    /// `xs`: `(B=1, T=1, H)`. Returns `(B=1, T=1, H)`.
    fn forward_step(
        &self,
        xs: &Tensor,
        cache_k: &mut Option<Tensor>,
        cache_v: &mut Option<Tensor>,
        rope: &RotaryCache,
        position: usize,
    ) -> Result<Tensor> {
        let (b, t, _h) = xs.dims3()?;
        debug_assert_eq!(t, 1, "self-attn step expects T=1");
        let q = self.q_proj.forward(xs)?;
        let k = self.k_proj.forward(xs)?;
        let v = self.v_proj.forward(xs)?;
        // (B, T, H) -> (B, H, T, head_dim)
        let q = q
            .reshape((b, t, self.n_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?;
        let k = k
            .reshape((b, t, self.n_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?;
        let v = v
            .reshape((b, t, self.n_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?;
        // RoPE for q and the new k at this position.
        let q = rope.apply_at(&q, position, 1)?;
        let k_new = rope.apply_at(&k, position, 1)?;
        // Collapse to (B*H, T, head_dim) for the matmul.
        let bh = b * self.n_heads;
        let q3 = q.reshape((bh, 1, self.head_dim))?;
        let k_new3 = k_new.reshape((bh, 1, self.head_dim))?;
        let v3 = v.reshape((bh, 1, self.head_dim))?;
        // Append to cache (or initialize on first step).
        let k_full = match cache_k.take() {
            Some(prev) => Tensor::cat(&[&prev, &k_new3], 1)?,
            None => k_new3,
        };
        let v_full = match cache_v.take() {
            Some(prev) => Tensor::cat(&[&prev, &v3], 1)?,
            None => v3,
        };
        // No causal mask needed: K/V only contains positions <= current.
        let scale = 1.0 / (self.head_dim as f64).sqrt();
        let scores = (q3.matmul(&k_full.transpose(1, 2)?.contiguous()?)? * scale)?;
        let probs = candle_nn::ops::softmax_last_dim(&scores)?;
        let out = probs.matmul(&v_full)?;
        // Write cache back.
        *cache_k = Some(k_full);
        *cache_v = Some(v_full);
        // (B*H, 1, head_dim) -> (B, 1, H)
        let out = out
            .reshape((b, self.n_heads, 1, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?
            .reshape((b, 1, self.n_heads * self.head_dim))?;
        self.o_proj.forward(&out)
    }
}

impl CrossAttention {
    /// Pre-compute and cache encoder K/V projections. Returns
    /// `(K, V)` each shaped `(B*H, T_enc, head_dim)`.
    fn precompute_kv(&self, enc: &Tensor) -> Result<(Tensor, Tensor)> {
        let (b, t_enc, _) = enc.dims3()?;
        let k = self.k_proj.forward(enc)?;
        let v = self.v_proj.forward(enc)?;
        let bh = b * self.n_heads;
        let k = k
            .reshape((b, t_enc, self.n_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?
            .reshape((bh, t_enc, self.head_dim))?;
        let v = v
            .reshape((b, t_enc, self.n_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?
            .reshape((bh, t_enc, self.head_dim))?;
        Ok((k, v))
    }

    /// Cached cross-attention step. K/V come from `precompute_kv`.
    /// `xs`: `(B=1, T=1, H)`. Returns `(B=1, T=1, H)`.
    fn forward_step(&self, xs: &Tensor, k: &Tensor, v: &Tensor) -> Result<Tensor> {
        let (b, t_dec, _) = xs.dims3()?;
        debug_assert_eq!(t_dec, 1, "cross-attn step expects T=1");
        let q = self.q_proj.forward(xs)?;
        let bh = b * self.n_heads;
        let q = q
            .reshape((b, t_dec, self.n_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?
            .reshape((bh, 1, self.head_dim))?;
        let scale = 1.0 / (self.head_dim as f64).sqrt();
        let scores = (q.matmul(&k.transpose(1, 2)?.contiguous()?)? * scale)?;
        let probs = candle_nn::ops::softmax_last_dim(&scores)?;
        let out = probs.matmul(v)?;
        let out = out
            .reshape((b, self.n_heads, 1, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?
            .reshape((b, 1, self.n_heads * self.head_dim))?;
        self.o_proj.forward(&out)
    }
}

impl Decoder {
    /// Pre-compute the encoder cross-attn K/V for every layer. Returns
    /// an empty `DecoderCache` ready for `step()` calls.
    pub fn precompute_cross_kv(&self, encoder_output: &Tensor) -> Result<DecoderCache> {
        let n = self.layers.len();
        let mut cross_k = Vec::with_capacity(n);
        let mut cross_v = Vec::with_capacity(n);
        for layer in &self.layers {
            let (k, v) = layer.cross_attn.precompute_kv(encoder_output)?;
            cross_k.push(k);
            cross_v.push(v);
        }
        Ok(DecoderCache {
            self_k: (0..n).map(|_| None).collect(),
            self_v: (0..n).map(|_| None).collect(),
            cross_k,
            cross_v,
            position: 0,
        })
    }

    /// Single-token cached forward. `token_id` is the *new* token to
    /// process. Returns logits for that one position, shape
    /// `(1, vocab)`. Updates `cache.position` and per-layer `self_k`/
    /// `self_v` in place.
    pub fn step(&self, token_id: u32, cache: &mut DecoderCache) -> Result<Tensor> {
        let device = self.embed_weight.device();
        let token = Tensor::from_vec(vec![token_id], (1, 1), device)?;
        let mut h = self.embed.forward(&token)?;
        let position = cache.position;
        for (i, layer) in self.layers.iter().enumerate() {
            let normed = layer.input_ln.forward(&h)?;
            let attn_out = layer.self_attn.forward_step(
                &normed,
                &mut cache.self_k[i],
                &mut cache.self_v[i],
                &self.rope,
                position,
            )?;
            h = (h + attn_out)?;
            let normed = layer.post_attn_ln.forward(&h)?;
            let cross_out =
                layer
                    .cross_attn
                    .forward_step(&normed, &cache.cross_k[i], &cache.cross_v[i])?;
            h = (h + cross_out)?;
            let normed = layer.final_ln.forward(&h)?;
            let mlp_out = layer.mlp.forward(&normed)?;
            h = (h + mlp_out)?;
        }
        h = self.final_ln.forward(&h)?;
        cache.position += 1;
        // Logits: (1, 1, H) -> (1, H) -> (1, vocab) via tied LM head.
        let lm_w = self.embed_weight.transpose(0, 1)?.contiguous()?;
        let logits_3d = h.broadcast_matmul(&lm_w)?;
        logits_3d.squeeze(1)
    }

    /// Cached greedy decode. Equivalent semantics to
    /// [`Self::generate`] but with O(T) total work instead of O(T²).
    pub fn generate_cached(
        &self,
        encoder_output: &Tensor,
        cfg: &MoonshineConfig,
        max_tokens: usize,
    ) -> Result<Vec<u32>> {
        let mut cache = self.precompute_cross_kv(encoder_output)?;
        let mut next = cfg.decoder_start_token_id;
        let mut out = Vec::with_capacity(max_tokens);
        for _ in 0..max_tokens {
            let logits = self.step(next, &mut cache)?; // (1, vocab)
            let argmax = logits.argmax(1)?;
            let id: u32 = argmax.to_dtype(DType::U32)?.to_vec1::<u32>()?[0];
            if id == cfg.eos_token_id {
                break;
            }
            out.push(id);
            next = id;
            if cache.position >= cfg.max_position_embeddings {
                break;
            }
        }
        Ok(out)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn config_padded_head_dim_pads_to_multiple_of_8() {
        let cfg = MoonshineConfig::tiny();
        // 288 / 8 = 36, padded up to 40 (next multiple of 8).
        assert_eq!(cfg.padded_head_dim(cfg.encoder_num_attention_heads), 40);
        // Sanity: smaller-head config still pads.
        let mut alt = cfg.clone();
        alt.hidden_size = 256;
        assert_eq!(alt.padded_head_dim(8), 32);
    }
}