inferencelayer 0.2.10

//! Qwen3-TTS-12Hz (Alibaba, Apache 2.0) — Phase 0 scaffolding.
//!
//! Zero-shot voice-cloning / preset-speaker / voice-design TTS. See `HANDOFF_qwen3tts.md` for the
//! full architecture, the staged port plan, and the **Rust-only gating discipline** (no torch
//! oracle; lineage-inherited parity + self-consistency gates + an end-to-end speaker-similarity
//! acceptance gate — the same three layers as `HANDOFF_chatterbox.md` §2).
//!
//! Pipeline (one 80 ms frame per AR step, 12.5 Hz, 16 codebooks × 2048):
//! ```text
//! text ─ Qwen3 BPE (vocab.json+merges.txt) ─┐            x-vector ─ ECAPA-TDNN (in-checkpoint) ─┐
//!   Talker = Qwen3 decoder (MRoPE, GQA+qk-norm), input = text-side ⊕ codec-side embed SUM ◄─────┘
//!     └─ 1 semantic code/frame ─► CodePredictor (5L MTP, fresh KV/frame) ─ 15 acoustic codes
//!         └─ 16 codes ─► codec decoder (RVQ dequant → win-72 transformer → ConvNeXt↑ → SnakeBeta
//!                          SEANet ×[8,5,4,3]) ─► 1920 samples @ 24 kHz
//! ```
//!
//! **Phase 0 (this file, for now):** config landmarks asserted against the checkpoint's own
//! `config.json`/`generation_config.json`, plus checkpoint-dir discovery. The manifest tooling is
//! shared with the chatterbox port (`chatterbox::{read_manifest, verify_shapes}`). Loaders and
//! forwards land in P1+ once `qwen3tts-dbg` confirms the ground-truth tensor manifest.

use crate::cpu_gemm::PackedWeight;
use anyhow::{Context, Result};
use serde_json::Value;
use std::path::{Path, PathBuf};

/// Config landmarks, quoted from the checkpoints' `config.json` (verified 2026-07-20 against
/// `Qwen/Qwen3-TTS-12Hz-1.7B-Base` + its bundled `speech_tokenizer/`, transformers 4.57.3).
/// The serde parse below is the loader's input; these constants are what the parse is ASSERTED
/// against for the recognized `tts_model_size` values, so a layout drift fails loudly at load.
pub mod cfg {
    // ---- Talker: the text→semantic-code Qwen3 decoder (`talker_config`) ----
    /// 1.7B ("1b7") hidden; 0.6B ("0b6") is 1024. Attention width is heads×head_dim = 2048 on BOTH
    /// (head_dim is pinned at 128 — NEVER derive it as hidden/heads; landmine §5.2).
    pub const TALKER_HIDDEN_1B7: usize = 2048;
    pub const TALKER_HIDDEN_0B6: usize = 1024;
    pub const TALKER_LAYERS: usize = 28;
    pub const TALKER_HEADS: usize = 16;
    pub const TALKER_KV_HEADS: usize = 8;
    pub const HEAD_DIM: usize = 128;
    pub const TALKER_INTERMEDIATE_1B7: usize = 6144;
    pub const TALKER_INTERMEDIATE_0B6: usize = 3072;
    pub const TALKER_RMS_EPS: f32 = 1e-6;
    pub const TALKER_ROPE_THETA: f32 = 1_000_000.0;
    /// Interleaved MRoPE over half-dim 64 = 24+20+20 (t,h,w). Positions are uniform across the
    /// three channels for this model — asserted at runtime, and uniform-MRoPE ≡ 1-D RoPE is an
    /// engine-pinned identity (`forward.rs` mrope tests).
    pub const MROPE_SECTION: [usize; 3] = [24, 20, 20];
    /// `position_id_per_seconds` from config — 13, NOT the 12.5 Hz frame rate. Its exact use site
    /// must be transcribed from the reference before P6 (landmine §5.1).
    pub const POSITION_ID_PER_SECONDS: usize = 13;

    // ---- Codec-token id space of the TALKER (vocab 3072 = 2048 codes + specials at 2048+) ----
    pub const TALKER_CODEC_VOCAB: usize = 3072;
    pub const CODEC_PAD: u32 = 2148;
    pub const CODEC_BOS: u32 = 2149;
    /// AR stop: the sampled SEMANTIC code equals this (CodePredictor is not run on that frame).
    pub const CODEC_EOS: u32 = 2150;
    pub const CODEC_THINK: u32 = 2154;
    pub const CODEC_NOTHINK: u32 = 2155;
    pub const CODEC_THINK_BOS: u32 = 2156;
    pub const CODEC_THINK_EOS: u32 = 2157;

    // ---- Text-side id space (Qwen3 BPE table, 151936 rows) ----
    pub const TEXT_VOCAB: usize = 151_936;
    pub const IM_START: u32 = 151_644;
    pub const IM_END: u32 = 151_645;
    pub const ASSISTANT: u32 = 77_091;
    pub const TTS_PAD: u32 = 151_671;
    pub const TTS_BOS: u32 = 151_672;
    pub const TTS_EOS: u32 = 151_673;

    // ---- CodePredictor (`talker_config.code_predictor_config`) — same dims on 1.7B and 0.6B ----
    pub const CP_HIDDEN: usize = 1024;
    pub const CP_LAYERS: usize = 5;
    pub const CP_HEADS: usize = 16;
    pub const CP_KV_HEADS: usize = 8;
    pub const CP_INTERMEDIATE: usize = 3072;
    /// Acoustic-code vocab (no specials — raw codes only; the THIRD id space, landmine §5.3).
    pub const CP_VOCAB: usize = 2048;
    /// 16 codebooks total: 1 semantic (talker) + 15 acoustic (code predictor).
    pub const NUM_CODE_GROUPS: usize = 16;
    pub const CP_ROPE_THETA: f32 = 1_000_000.0;

    // ---- Frame geometry (identical to Mimi's) ----
    pub const SAMPLE_RATE: u32 = 24_000;
    /// `decode_upsample_rate`/`encode_downsample_rate`: samples per 12.5 Hz frame.
    pub const SAMPLES_PER_FRAME: usize = 1920;
    /// generation_config `max_new_tokens`: AR frame cap (~10.9 min).
    pub const MAX_FRAMES: usize = 8192;

    // ---- Codec (bundled `speech_tokenizer/`, = `Qwen/Qwen3-TTS-Tokenizer-12Hz`) ----
    /// Decoder: RVQ codebook_dim 512, hidden 512, 8 layers, sliding window 72, LayerScale 0.01,
    /// upsampling_ratios [2,2] (ConvTr+ConvNeXt, 12.5→50 Hz) then SEANet blocks ×[8,5,4,3].
    pub const DEC_HIDDEN: usize = 512;
    pub const DEC_LAYERS: usize = 8;
    pub const DEC_HEADS: usize = 16; // MHA (kv_heads == heads), head_dim 64
    pub const DEC_HEAD_DIM: usize = 64;
    pub const DEC_LATENT: usize = 1024;
    pub const DEC_CODEBOOK_DIM: usize = 512;
    pub const DEC_DECODER_DIM: usize = 1536;
    pub const DEC_SLIDING_WINDOW: usize = 72;
    pub const DEC_UPSAMPLE_RATES: [usize; 4] = [8, 5, 4, 3];
    pub const DEC_RMS_EPS: f32 = 1e-5;
    pub const DEC_ROPE_THETA: f32 = 10_000.0;
    /// `semantic_codebook_size: 4096` vs codebook_size 2048 — resolved from the manifest at P0
    /// before the RVQ loader is written (landmine §5.4).
    pub const DEC_SEMANTIC_CODEBOOK_SIZE: usize = 4096;
    /// Encoder: Mimi-lineage SEANet (ratios 8·6·5·4 = 960× → 25 Hz, then stride-2 → 12.5 Hz),
    /// 8-layer GELU transformer, sliding window 250, 32 quantizers of which 16 are valid.
    pub const ENC_HIDDEN: usize = 512;
    pub const ENC_LAYERS: usize = 8;
    pub const ENC_HEADS: usize = 8;
    pub const ENC_CODEBOOK_DIM: usize = 256;
    pub const ENC_QUANTIZERS_TOTAL: usize = 32;
    pub const ENC_QUANTIZERS_VALID: usize = 16;
    pub const ENC_SLIDING_WINDOW: usize = 250;
    pub const ENC_RATIOS: [usize; 4] = [8, 6, 5, 4];

    // ---- Sampling (generation_config.json; subtalker = CodePredictor) ----
    pub const SAMPLE_REPETITION_PENALTY: f32 = 1.05; // talker only — subtalker has NO rep-pen
    pub const SAMPLE_TEMPERATURE: f32 = 0.9;
    pub const SAMPLE_TOP_K: usize = 50;
    pub const SAMPLE_TOP_P: f32 = 1.0;

    // ---- Checkpoint layout (one HF snapshot dir per variant) ----
    pub const FILE_MODEL: &str = "model.safetensors";
    pub const FILE_CONFIG: &str = "config.json";
    pub const FILE_GENERATION: &str = "generation_config.json";
    pub const FILE_VOCAB: &str = "vocab.json";
    pub const FILE_MERGES: &str = "merges.txt";
    pub const DIR_CODEC: &str = "speech_tokenizer";
}

/// `~/.cache/inferencelayer/qwen3tts/<variant>` (or `$OSFKB_QWEN3TTS_DIR/<variant>`); the tests
/// and the dbg tool look for `1.7b-base` first. Mirrors the chatterbox cache convention.
pub fn default_root() -> PathBuf {
    std::env::var("OSFKB_QWEN3TTS_DIR")
        .map(PathBuf::from)
        .unwrap_or_else(|_| {
            PathBuf::from(std::env::var("HOME").unwrap_or_else(|_| ".".into()))
                .join(".cache/inferencelayer/qwen3tts")
        })
}

// ---------------------------------------------------------------------------------------------
// config.json parse — the checkpoint's own metadata, asserted against `cfg` at load.
// Manual serde_json::Value extraction (the crate's serde-derive feature is off the encoder-cpu
// build; this is the house pattern — see kokoro.rs / moshi_lm.rs).
// ---------------------------------------------------------------------------------------------

fn get<'a>(v: &'a Value, k: &str) -> Result<&'a Value> {
    v.get(k)
        .with_context(|| format!("missing config key {k:?}"))
}
fn ju(v: &Value, k: &str) -> Result<usize> {
    Ok(get(v, k)?
        .as_u64()
        .with_context(|| format!("config key {k:?} is not an integer"))? as usize)
}
fn ju32(v: &Value, k: &str) -> Result<u32> {
    Ok(get(v, k)?
        .as_u64()
        .with_context(|| format!("config key {k:?} is not an integer"))? as u32)
}
fn jf(v: &Value, k: &str) -> Result<f32> {
    Ok(get(v, k)?
        .as_f64()
        .with_context(|| format!("config key {k:?} is not a number"))? as f32)
}
fn js(v: &Value, k: &str) -> Result<String> {
    Ok(get(v, k)?
        .as_str()
        .with_context(|| format!("config key {k:?} is not a string"))?
        .to_string())
}
fn jb(v: &Value, k: &str) -> Result<bool> {
    get(v, k)?
        .as_bool()
        .with_context(|| format!("config key {k:?} is not a bool"))
}
fn juvec(v: &Value, k: &str) -> Result<Vec<usize>> {
    get(v, k)?
        .as_array()
        .with_context(|| format!("config key {k:?} is not an array"))?
        .iter()
        .map(|x| {
            x.as_u64()
                .map(|x| x as usize)
                .context("non-integer array element")
        })
        .collect()
}
/// Optional `{name: id}` map (absent or empty on some variants).
fn jmap_u32(v: &Value, k: &str) -> std::collections::BTreeMap<String, u32> {
    let mut out = std::collections::BTreeMap::new();
    if let Some(m) = v.get(k).and_then(|m| m.as_object()) {
        for (name, id) in m {
            if let Some(id) = id.as_u64() {
                out.insert(name.clone(), id as u32);
            }
        }
    }
    out
}

#[derive(Debug, Clone)]
pub struct TopConfig {
    pub tts_model_size: String, // "1b7" | "0b6"
    pub tts_model_type: String, // "base" | "custom_voice" | "voice_design" (spelling per checkpoint)
    pub speaker_encoder_config: Option<SpkEncConfig>,
    pub talker_config: TalkerConfig,
}

#[derive(Debug, Clone)]
pub struct SpkEncConfig {
    pub enc_dim: usize,
    pub sample_rate: u32,
}

#[derive(Debug, Clone)]
pub struct TalkerConfig {
    pub hidden_size: usize,
    pub num_hidden_layers: usize,
    pub num_attention_heads: usize,
    pub num_key_value_heads: usize,
    pub head_dim: usize,
    pub intermediate_size: usize,
    pub rms_norm_eps: f32,
    pub rope_theta: f32,
    pub mrope_section: [usize; 3],
    pub mrope_interleaved: bool,
    pub vocab_size: usize,      // codec-side table: 3072
    pub text_vocab_size: usize, // text-side table: 151936
    pub text_hidden_size: usize,
    pub num_code_groups: usize,
    pub position_id_per_seconds: usize,
    pub code_predictor_config: CodePredictorConfig,
    pub codec_bos_id: u32,
    pub codec_eos_id: u32,
    pub codec_pad_id: u32,
    pub codec_think_id: u32,
    pub codec_nothink_id: u32,
    pub codec_think_bos_id: u32,
    pub codec_think_eos_id: u32,
    /// language name → codec-side id (2050..=2071).
    pub codec_language_id: std::collections::BTreeMap<String, u32>,
    /// CustomVoice/VoiceDesign: speaker name → codec-side embedding row. Empty on Base.
    pub spk_id: std::collections::BTreeMap<String, u32>,
    /// CustomVoice: speaker name → dialect language key (`Some("sichuan_dialect")`) or None
    /// (config value `false`). Drives the language-id override for chinese/auto requests.
    pub spk_is_dialect: std::collections::BTreeMap<String, Option<String>>,
}

#[derive(Debug, Clone)]
pub struct CodePredictorConfig {
    pub hidden_size: usize,
    pub num_hidden_layers: usize,
    pub num_attention_heads: usize,
    pub num_key_value_heads: usize,
    pub head_dim: usize,
    pub intermediate_size: usize,
    pub rms_norm_eps: f32,
    pub rope_theta: f32,
    pub vocab_size: usize,
    pub num_code_groups: usize,
}

/// `speech_tokenizer/config.json` (== the standalone `Qwen3-TTS-Tokenizer-12Hz`).
#[derive(Debug, Clone)]
pub struct CodecConfig {
    pub encoder_valid_num_quantizers: usize,
    pub input_sample_rate: u32,
    pub output_sample_rate: u32,
    pub decode_upsample_rate: usize,
    pub encode_downsample_rate: usize,
    pub decoder_config: CodecDecoderConfig,
    pub encoder_config: CodecEncoderConfig,
}

#[derive(Debug, Clone)]
pub struct CodecDecoderConfig {
    pub latent_dim: usize,
    pub codebook_dim: usize,
    pub codebook_size: usize,
    pub decoder_dim: usize,
    pub hidden_size: usize,
    pub intermediate_size: usize,
    pub layer_scale_initial_scale: f32,
    pub head_dim: usize,
    pub num_attention_heads: usize,
    pub num_hidden_layers: usize,
    pub num_key_value_heads: usize,
    pub num_quantizers: usize,
    pub num_semantic_quantizers: usize,
    pub rms_norm_eps: f32,
    pub rope_theta: f32,
    pub semantic_codebook_size: usize,
    pub sliding_window: usize,
    pub upsample_rates: Vec<usize>,
    pub upsampling_ratios: Vec<usize>,
}

#[derive(Debug, Clone)]
pub struct CodecEncoderConfig {
    pub codebook_dim: usize,
    pub codebook_size: usize,
    pub hidden_size: usize,
    pub intermediate_size: usize,
    pub kernel_size: usize,
    pub last_kernel_size: usize,
    pub num_attention_heads: usize,
    pub num_filters: usize,
    pub num_hidden_layers: usize,
    pub num_key_value_heads: usize,
    pub num_quantizers: usize,
    pub num_residual_layers: usize,
    pub num_semantic_quantizers: usize,
    pub residual_kernel_size: usize,
    pub sliding_window: usize,
    pub trim_right_ratio: f32,
    pub upsampling_ratios: Vec<usize>,
    pub use_causal_conv: bool,
}

/// generation_config.json — the checkpoint's own sampling defaults (subtalker = CodePredictor).
#[derive(Debug, Clone)]
pub struct GenerationConfig {
    pub do_sample: bool,
    pub repetition_penalty: f32,
    pub temperature: f32,
    pub top_p: f32,
    pub top_k: usize,
    pub subtalker_temperature: f32,
    pub subtalker_top_p: f32,
    pub subtalker_top_k: usize,
    pub max_new_tokens: usize,
}

impl TopConfig {
    fn from_json(v: &Value) -> Result<Self> {
        let spk = v
            .get("speaker_encoder_config")
            .filter(|s| !s.is_null())
            .map(|s| -> Result<SpkEncConfig> {
                Ok(SpkEncConfig {
                    enc_dim: ju(s, "enc_dim")?,
                    sample_rate: ju32(s, "sample_rate")?,
                })
            })
            .transpose()?;
        Ok(Self {
            tts_model_size: js(v, "tts_model_size")?,
            tts_model_type: js(v, "tts_model_type")?,
            speaker_encoder_config: spk,
            talker_config: TalkerConfig::from_json(get(v, "talker_config")?)?,
        })
    }
}

impl TalkerConfig {
    fn from_json(v: &Value) -> Result<Self> {
        let rs = get(v, "rope_scaling")?;
        let sect = juvec(rs, "mrope_section")?;
        anyhow::ensure!(
            sect.len() == 3,
            "mrope_section has {} entries, want 3",
            sect.len()
        );
        Ok(Self {
            hidden_size: ju(v, "hidden_size")?,
            num_hidden_layers: ju(v, "num_hidden_layers")?,
            num_attention_heads: ju(v, "num_attention_heads")?,
            num_key_value_heads: ju(v, "num_key_value_heads")?,
            head_dim: ju(v, "head_dim")?,
            intermediate_size: ju(v, "intermediate_size")?,
            rms_norm_eps: jf(v, "rms_norm_eps")?,
            rope_theta: jf(v, "rope_theta")?,
            mrope_section: [sect[0], sect[1], sect[2]],
            mrope_interleaved: rs
                .get("interleaved")
                .and_then(|b| b.as_bool())
                .unwrap_or(false),
            vocab_size: ju(v, "vocab_size")?,
            text_vocab_size: ju(v, "text_vocab_size")?,
            text_hidden_size: ju(v, "text_hidden_size")?,
            num_code_groups: ju(v, "num_code_groups")?,
            position_id_per_seconds: ju(v, "position_id_per_seconds")?,
            code_predictor_config: CodePredictorConfig::from_json(get(
                v,
                "code_predictor_config",
            )?)?,
            codec_bos_id: ju32(v, "codec_bos_id")?,
            codec_eos_id: ju32(v, "codec_eos_token_id")?,
            codec_pad_id: ju32(v, "codec_pad_id")?,
            codec_think_id: ju32(v, "codec_think_id")?,
            codec_nothink_id: ju32(v, "codec_nothink_id")?,
            codec_think_bos_id: ju32(v, "codec_think_bos_id")?,
            codec_think_eos_id: ju32(v, "codec_think_eos_id")?,
            codec_language_id: jmap_u32(v, "codec_language_id"),
            spk_id: jmap_u32(v, "spk_id"),
            spk_is_dialect: {
                let mut out = std::collections::BTreeMap::new();
                if let Some(m) = v.get("spk_is_dialect").and_then(|m| m.as_object()) {
                    for (name, val) in m {
                        out.insert(name.clone(), val.as_str().map(|s| s.to_string()));
                    }
                }
                out
            },
        })
    }
}

impl CodePredictorConfig {
    fn from_json(v: &Value) -> Result<Self> {
        Ok(Self {
            hidden_size: ju(v, "hidden_size")?,
            num_hidden_layers: ju(v, "num_hidden_layers")?,
            num_attention_heads: ju(v, "num_attention_heads")?,
            num_key_value_heads: ju(v, "num_key_value_heads")?,
            head_dim: ju(v, "head_dim")?,
            intermediate_size: ju(v, "intermediate_size")?,
            rms_norm_eps: jf(v, "rms_norm_eps")?,
            rope_theta: jf(v, "rope_theta")?,
            vocab_size: ju(v, "vocab_size")?,
            num_code_groups: ju(v, "num_code_groups")?,
        })
    }
}

impl CodecConfig {
    fn from_json(v: &Value) -> Result<Self> {
        Ok(Self {
            encoder_valid_num_quantizers: ju(v, "encoder_valid_num_quantizers")?,
            input_sample_rate: ju32(v, "input_sample_rate")?,
            output_sample_rate: ju32(v, "output_sample_rate")?,
            decode_upsample_rate: ju(v, "decode_upsample_rate")?,
            encode_downsample_rate: ju(v, "encode_downsample_rate")?,
            decoder_config: CodecDecoderConfig::from_json(get(v, "decoder_config")?)?,
            encoder_config: CodecEncoderConfig::from_json(get(v, "encoder_config")?)?,
        })
    }
}

impl CodecDecoderConfig {
    fn from_json(v: &Value) -> Result<Self> {
        Ok(Self {
            latent_dim: ju(v, "latent_dim")?,
            codebook_dim: ju(v, "codebook_dim")?,
            codebook_size: ju(v, "codebook_size")?,
            decoder_dim: ju(v, "decoder_dim")?,
            hidden_size: ju(v, "hidden_size")?,
            intermediate_size: ju(v, "intermediate_size")?,
            layer_scale_initial_scale: jf(v, "layer_scale_initial_scale")?,
            head_dim: ju(v, "head_dim")?,
            num_attention_heads: ju(v, "num_attention_heads")?,
            num_hidden_layers: ju(v, "num_hidden_layers")?,
            num_key_value_heads: ju(v, "num_key_value_heads")?,
            num_quantizers: ju(v, "num_quantizers")?,
            num_semantic_quantizers: ju(v, "num_semantic_quantizers")?,
            rms_norm_eps: jf(v, "rms_norm_eps")?,
            rope_theta: jf(v, "rope_theta")?,
            semantic_codebook_size: ju(v, "semantic_codebook_size")?,
            sliding_window: ju(v, "sliding_window")?,
            upsample_rates: juvec(v, "upsample_rates")?,
            upsampling_ratios: juvec(v, "upsampling_ratios")?,
        })
    }
}

impl CodecEncoderConfig {
    fn from_json(v: &Value) -> Result<Self> {
        Ok(Self {
            codebook_dim: ju(v, "codebook_dim")?,
            codebook_size: ju(v, "codebook_size")?,
            hidden_size: ju(v, "hidden_size")?,
            intermediate_size: ju(v, "intermediate_size")?,
            kernel_size: ju(v, "kernel_size")?,
            last_kernel_size: ju(v, "last_kernel_size")?,
            num_attention_heads: ju(v, "num_attention_heads")?,
            num_filters: ju(v, "num_filters")?,
            num_hidden_layers: ju(v, "num_hidden_layers")?,
            num_key_value_heads: ju(v, "num_key_value_heads")?,
            num_quantizers: ju(v, "num_quantizers")?,
            num_residual_layers: ju(v, "num_residual_layers")?,
            num_semantic_quantizers: ju(v, "num_semantic_quantizers")?,
            residual_kernel_size: ju(v, "residual_kernel_size")?,
            sliding_window: ju(v, "sliding_window")?,
            trim_right_ratio: jf(v, "trim_right_ratio")?,
            upsampling_ratios: juvec(v, "upsampling_ratios")?,
            use_causal_conv: jb(v, "use_causal_conv")?,
        })
    }
}

impl GenerationConfig {
    fn from_json(v: &Value) -> Result<Self> {
        Ok(Self {
            do_sample: jb(v, "do_sample")?,
            repetition_penalty: jf(v, "repetition_penalty")?,
            temperature: jf(v, "temperature")?,
            top_p: jf(v, "top_p")?,
            top_k: ju(v, "top_k")?,
            subtalker_temperature: jf(v, "subtalker_temperature")?,
            subtalker_top_p: jf(v, "subtalker_top_p")?,
            subtalker_top_k: ju(v, "subtalker_top_k")?,
            max_new_tokens: ju(v, "max_new_tokens")?,
        })
    }
}

/// The parsed identity of one checkpoint dir (talker + bundled codec), validated against `cfg`.
#[derive(Debug, Clone)]
pub struct Qwen3TtsConfig {
    pub top: TopConfig,
    pub codec: CodecConfig,
    pub generation: GenerationConfig,
}

impl Qwen3TtsConfig {
    /// Parse `<dir>/{config,generation_config}.json` + `<dir>/speech_tokenizer/config.json` and
    /// fail loudly on any deviation from the landmarks in [`cfg`]. This is the P0 gate: a checkpoint
    /// that parses but mismatches would otherwise miscompute silently 10 stages later.
    pub fn load(dir: &Path) -> Result<Self> {
        let read = |p: PathBuf| -> Result<Value> {
            serde_json::from_slice(
                &std::fs::read(&p).with_context(|| format!("read {}", p.display()))?,
            )
            .with_context(|| format!("parse {}", p.display()))
        };
        let top = TopConfig::from_json(&read(dir.join(cfg::FILE_CONFIG))?)
            .context("parse config.json")?;
        let codec = CodecConfig::from_json(&read(dir.join(cfg::DIR_CODEC).join(cfg::FILE_CONFIG))?)
            .context("parse speech_tokenizer/config.json")?;
        let generation = GenerationConfig::from_json(&read(dir.join(cfg::FILE_GENERATION))?)
            .context("parse generation_config.json")?;
        let s = Self {
            top,
            codec,
            generation,
        };
        s.validate()?;
        Ok(s)
    }

    fn validate(&self) -> Result<()> {
        use anyhow::ensure;
        let t = &self.top.talker_config;
        let (want_h, want_i) = match self.top.tts_model_size.as_str() {
            "1b7" => (cfg::TALKER_HIDDEN_1B7, cfg::TALKER_INTERMEDIATE_1B7),
            "0b6" => (cfg::TALKER_HIDDEN_0B6, cfg::TALKER_INTERMEDIATE_0B6),
            other => anyhow::bail!("unrecognized tts_model_size {other:?} (want 1b7|0b6)"),
        };
        ensure!(
            t.hidden_size == want_h,
            "talker hidden {} ≠ {want_h}",
            t.hidden_size
        );
        ensure!(
            t.intermediate_size == want_i,
            "talker intermediate {}",
            t.intermediate_size
        );
        ensure!(
            t.num_hidden_layers == cfg::TALKER_LAYERS,
            "talker layers {}",
            t.num_hidden_layers
        );
        ensure!(t.num_attention_heads == cfg::TALKER_HEADS, "talker heads");
        ensure!(
            t.num_key_value_heads == cfg::TALKER_KV_HEADS,
            "talker kv heads"
        );
        ensure!(
            t.head_dim == cfg::HEAD_DIM,
            "talker head_dim {} ≠ 128",
            t.head_dim
        );
        ensure!(t.mrope_section == cfg::MROPE_SECTION, "mrope_section");
        ensure!(t.mrope_interleaved, "expected interleaved mrope");
        ensure!(
            t.vocab_size == cfg::TALKER_CODEC_VOCAB,
            "talker codec vocab {}",
            t.vocab_size
        );
        ensure!(
            t.text_vocab_size == cfg::TEXT_VOCAB,
            "text vocab {}",
            t.text_vocab_size
        );
        ensure!(t.num_code_groups == cfg::NUM_CODE_GROUPS, "num_code_groups");
        ensure!(
            t.codec_eos_id == cfg::CODEC_EOS,
            "codec_eos {}",
            t.codec_eos_id
        );
        ensure!(
            t.codec_bos_id == cfg::CODEC_BOS && t.codec_pad_id == cfg::CODEC_PAD,
            "codec bos/pad"
        );

        let p = &t.code_predictor_config;
        ensure!(
            p.hidden_size == cfg::CP_HIDDEN,
            "cp hidden {}",
            p.hidden_size
        );
        ensure!(p.num_hidden_layers == cfg::CP_LAYERS, "cp layers");
        ensure!(
            p.head_dim == cfg::HEAD_DIM,
            "cp head_dim {} ≠ 128",
            p.head_dim
        );
        ensure!(p.vocab_size == cfg::CP_VOCAB, "cp vocab {}", p.vocab_size);
        ensure!(
            p.num_code_groups == cfg::NUM_CODE_GROUPS,
            "cp num_code_groups"
        );

        let d = &self.codec.decoder_config;
        ensure!(
            self.codec.decode_upsample_rate == cfg::SAMPLES_PER_FRAME,
            "frame size"
        );
        ensure!(
            self.codec.encode_downsample_rate == cfg::SAMPLES_PER_FRAME,
            "frame size (enc)"
        );
        ensure!(
            self.codec.output_sample_rate == cfg::SAMPLE_RATE,
            "sample rate"
        );
        ensure!(
            self.codec.encoder_valid_num_quantizers == cfg::ENC_QUANTIZERS_VALID,
            "valid quantizers"
        );
        ensure!(
            d.hidden_size == cfg::DEC_HIDDEN && d.num_hidden_layers == cfg::DEC_LAYERS,
            "dec core"
        );
        ensure!(
            d.sliding_window == cfg::DEC_SLIDING_WINDOW,
            "dec sliding window"
        );
        ensure!(
            d.codebook_size == cfg::CP_VOCAB,
            "dec codebook size {}",
            d.codebook_size
        );
        ensure!(
            d.upsample_rates == cfg::DEC_UPSAMPLE_RATES,
            "dec upsample rates"
        );
        ensure!(d.num_quantizers == cfg::NUM_CODE_GROUPS, "dec quantizers");

        let e = &self.codec.encoder_config;
        ensure!(e.upsampling_ratios == cfg::ENC_RATIOS, "enc ratios");
        ensure!(
            e.num_quantizers == cfg::ENC_QUANTIZERS_TOTAL,
            "enc quantizers"
        );
        ensure!(e.use_causal_conv, "enc must be causal");
        // The codec forwards use cfg constants for eps/theta/head_dim — assert the checkpoint
        // agrees so a future speech_tokenizer variant fails loudly instead of diverging silently
        // (review finding).
        ensure!(
            d.rms_norm_eps == cfg::DEC_RMS_EPS,
            "dec rms eps {}",
            d.rms_norm_eps
        );
        ensure!(
            d.rope_theta == cfg::DEC_ROPE_THETA,
            "dec rope theta {}",
            d.rope_theta
        );
        ensure!(
            d.head_dim == cfg::DEC_HEAD_DIM,
            "dec head_dim {}",
            d.head_dim
        );
        ensure!(
            d.num_attention_heads == cfg::DEC_HEADS && d.latent_dim == cfg::DEC_LATENT,
            "dec attn/latent dims"
        );
        ensure!(
            e.sliding_window == cfg::ENC_SLIDING_WINDOW,
            "enc sliding window"
        );
        ensure!(
            e.hidden_size == cfg::ENC_HIDDEN && e.codebook_dim == cfg::ENC_CODEBOOK_DIM,
            "enc dims"
        );
        Ok(())
    }
}

// ---------------------------------------------------------------------------------------------
// P1 — text tokenizer. The snapshot ships NO tokenizer.json, only vocab.json + merges.txt, so the
// pipeline is assembled programmatically to match the family's official tokenizer.json exactly
// (verified against Qwen/Qwen3-0.6B: NFC normalizer → Split(Qwen2 regex, Isolated) →
// ByteLevel(add_prefix_space=false, use_regex=false) → BPE → ByteLevel decoder). Added tokens
// (33, ids 151643..=151675, contiguous) come from tokenizer_config.json's `added_tokens_decoder`
// and are registered in ascending-id order so the crate assigns them their recorded ids.
// Landmine §5.12: the crate's ByteLevel DEFAULT regex is GPT-2's, not Qwen2's — hence the
// explicit Split stage and the pinned-id gates in tests/qwen3tts_parity.rs.
// ---------------------------------------------------------------------------------------------

#[cfg(feature = "cli")]
mod text_tokenizer {
    use super::cfg;
    use anyhow::{Context, Result};
    use std::path::Path;

    /// transformers `Qwen2Tokenizer` PRETOKENIZE_REGEX — byte-identical to the `Split` pattern
    /// recorded in the family's official tokenizer.json.
    pub const QWEN2_SPLIT_REGEX: &str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+";

    pub struct TextTokenizer {
        inner: tokenizers::Tokenizer,
    }

    impl TextTokenizer {
        pub fn load(dir: &Path) -> Result<Self> {
            use tokenizers::AddedToken;
            use tokenizers::SplitDelimiterBehavior;
            use tokenizers::models::bpe::BPE;
            use tokenizers::normalizers::unicode::NFC;
            use tokenizers::pre_tokenizers::byte_level::ByteLevel;
            use tokenizers::pre_tokenizers::sequence::Sequence;
            use tokenizers::pre_tokenizers::split::{Split, SplitPattern};

            let vocab = dir.join(cfg::FILE_VOCAB);
            let merges = dir.join(cfg::FILE_MERGES);
            let bpe = BPE::from_file(
                vocab.to_str().context("vocab path not utf-8")?,
                merges.to_str().context("merges path not utf-8")?,
            )
            .build()
            .map_err(|e| anyhow::anyhow!("build BPE from {}: {e}", dir.display()))?;

            let mut tk = tokenizers::Tokenizer::new(bpe);
            tk.with_normalizer(Some(NFC));
            let split = Split::new(
                SplitPattern::Regex(QWEN2_SPLIT_REGEX.into()),
                SplitDelimiterBehavior::Isolated,
                false,
            )
            .map_err(|e| anyhow::anyhow!("split pre-tokenizer: {e}"))?;
            // add_prefix_space=false, trim_offsets=false, use_regex=false — the GPT-2 regex inside
            // ByteLevel must stay OFF; the Split stage above is the only segmentation.
            let byte_level = ByteLevel::new(false, false, false);
            tk.with_pre_tokenizer(Some(Sequence::new(vec![split.into(), byte_level.into()])));
            tk.with_decoder(Some(tokenizers::decoders::byte_level::ByteLevel::new(
                false, false, false,
            )));

            // Added tokens, ascending id; assert each lands on its recorded id.
            let tc: serde_json::Value = serde_json::from_slice(
                &std::fs::read(dir.join("tokenizer_config.json"))
                    .with_context(|| format!("read tokenizer_config.json in {}", dir.display()))?,
            )
            .context("parse tokenizer_config.json")?;
            let decoder = tc
                .get("added_tokens_decoder")
                .and_then(|d| d.as_object())
                .context("tokenizer_config.json: missing added_tokens_decoder")?;
            let mut added: Vec<(u32, String, bool)> = Vec::with_capacity(decoder.len());
            for (id, meta) in decoder {
                let id: u32 = id.parse().context("added token id not an integer")?;
                let content = meta
                    .get("content")
                    .and_then(|c| c.as_str())
                    .context("added token without content")?
                    .to_string();
                let special = meta
                    .get("special")
                    .and_then(|s| s.as_bool())
                    .unwrap_or(false);
                added.push((id, content, special));
            }
            added.sort_by_key(|(id, _, _)| *id);
            for (id, content, special) in &added {
                let tok = AddedToken::from(content.clone(), *special);
                if *special {
                    tk.add_special_tokens(&[tok]);
                } else {
                    tk.add_tokens(&[tok]);
                }
                let got = tk.token_to_id(content);
                anyhow::ensure!(
                    got == Some(*id),
                    "added token {content:?} landed on id {got:?}, want {id} — \
                     added_tokens_decoder is not contiguous from the base vocab"
                );
            }
            Ok(Self { inner: tk })
        }

        /// Text → ids. Added/special tokens present in the text are matched (the reference
        /// behavior: `split_special_tokens=false`); no template tokens are inserted.
        pub fn encode(&self, text: &str) -> Result<Vec<u32>> {
            let enc = self
                .inner
                .encode(text, false)
                .map_err(|e| anyhow::anyhow!("encode: {e}"))?;
            Ok(enc.get_ids().to_vec())
        }

        /// Ids → text, specials kept verbatim.
        pub fn decode(&self, ids: &[u32]) -> Result<String> {
            self.inner
                .decode(ids, false)
                .map_err(|e| anyhow::anyhow!("decode: {e}"))
        }

        pub fn token_to_id(&self, token: &str) -> Option<u32> {
            self.inner.token_to_id(token)
        }

        /// Base vocab + added tokens (151643 + 33 = 151676; the embedding table is padded to
        /// 151936 — ids are always < the table height, never == it).
        pub fn vocab_size(&self) -> usize {
            self.inner.get_vocab_size(true)
        }
    }
}

#[cfg(feature = "cli")]
pub use text_tokenizer::{QWEN2_SPLIT_REGEX, TextTokenizer};

// ---------------------------------------------------------------------------------------------
// P2 — Talker backbone (CPU, the parity oracle for the P9 GPU offload).
//
// A Qwen3-style decoder with three deviations from the stock LLM spine that keep this a
// self-contained module (moshi/chatterbox precedent) rather than an Arch dispatch:
//   1. dual-stream input embeds — text-side `text_projection(text_embedding[id])` summed with
//      codec-side `codec_embedding[id]` per position (assembled by the P6 prompt builder; the
//      backbone itself takes ready-made `[seq, h]` embeds);
//   2. attention width ≠ hidden: q is 16×128=2048, kv 8×128=1024, GQA 2:1, per-head RMS qk-norm —
//      on the 0.6B (h=1024) and the CodePredictor these are WIDER than hidden (landmine §5.2);
//   3. interleaved MRoPE [24,20,20] over half=64 — positions are uniform (t=h=w cumsum; verified
//      in the reference `get_rope_index`), asserted here, and uniform-MRoPE ≡ 1-D RoPE is an
//      engine-pinned identity (gate `p2_mrope_uniform_equals_1d`).
//
// Two forward paths on purpose: `forward_hidden` (whole-sequence, no cache — the reference) and
// `prefill`/`step` over a `TalkerState` KV cache (the product). Every op accumulates in the same
// order in both, so the gate `p2_step_equals_full_forward` demands BIT-EXACT equality.
// ---------------------------------------------------------------------------------------------

pub(crate) fn matvec(m: &[f32], x: &[f32], rows: usize, cols: usize) -> Vec<f32> {
    debug_assert_eq!(m.len(), rows * cols);
    debug_assert_eq!(x.len(), cols);
    (0..rows)
        .map(|r| {
            m[r * cols..(r + 1) * cols]
                .iter()
                .zip(x)
                .map(|(a, b)| a * b)
                .sum()
        })
        .collect()
}

pub(crate) fn rms_norm(x: &[f32], w: &[f32], eps: f32) -> Vec<f32> {
    let ms = x.iter().map(|v| v * v).sum::<f32>() / x.len() as f32;
    let inv = 1.0 / (ms + eps).sqrt();
    x.iter().zip(w).map(|(v, g)| v * inv * g).collect()
}

pub(crate) fn softmax_inplace(v: &mut [f32]) {
    let m = v.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let mut s = 0.0;
    for x in v.iter_mut() {
        *x = (*x - m).exp();
        s += *x;
    }
    for x in v.iter_mut() {
        *x /= s;
    }
}

pub(crate) fn silu(x: f32) -> f32 {
    x / (1.0 + (-x).exp())
}

/// Which axis (0=t, 1=h, 2=w) each rotary frequency rotates by under INTERLEAVED M-RoPE — the
/// same derivation as the engine's `forward.rs::mrope_axis_map` (private there; copied with
/// provenance, non-chunked branch only). For `[24,20,20]` over half=64: h claims 1,4,…,58,
/// w claims 2,5,…,59, everything else (incl. 60..64) stays t.
fn mrope_axis_map(section: [usize; 3], half: usize) -> Vec<u8> {
    let mut axes = vec![0u8; half];
    for (d, &len) in section.iter().enumerate().skip(1) {
        let mut i = d;
        while i < (len * 3).min(half) {
            axes[i] = d as u8;
            i += 3;
        }
    }
    axes
}

/// One position's rope cos/sin over `hd` dims (`forward.rs::mrope_cs_padded`, full-rotary case).
/// With uniform `pos3` this is bitwise the 1-D table — gate `p2_mrope_uniform_equals_1d`.
fn mrope_cs(hd: usize, theta: f32, axes: &[u8], pos3: [u32; 3]) -> (Vec<f32>, Vec<f32>) {
    let half = hd / 2;
    debug_assert_eq!(axes.len(), half);
    let (mut cos, mut sin) = (vec![0f32; hd], vec![0f32; hd]);
    for i in 0..half {
        let pos = pos3[axes[i] as usize];
        let f = (pos as f32) * theta.powf(-2.0 * i as f32 / hd as f32);
        let (s, c) = f.sin_cos();
        cos[i] = c;
        cos[i + half] = c;
        sin[i] = s;
        sin[i + half] = s;
    }
    (cos, sin)
}

/// 1-D rope table, same expression as `mrope_cs` — exists only for the identity gate.
pub(crate) fn rope_cs_1d(hd: usize, theta: f32, pos: u32) -> (Vec<f32>, Vec<f32>) {
    let half = hd / 2;
    let (mut cos, mut sin) = (vec![0f32; hd], vec![0f32; hd]);
    for i in 0..half {
        let f = (pos as f32) * theta.powf(-2.0 * i as f32 / hd as f32);
        let (s, c) = f.sin_cos();
        cos[i] = c;
        cos[i + half] = c;
        sin[i] = s;
        sin[i + half] = s;
    }
    (cos, sin)
}

/// HF `rotate_half` rope applied in place to `heads × hd` at one position.
pub(crate) fn apply_rope(x: &mut [f32], cos: &[f32], sin: &[f32], heads: usize, hd: usize) {
    let half = hd / 2;
    for head in 0..heads {
        let base = head * hd;
        let mut rot = vec![0f32; hd];
        for d in 0..half {
            rot[d] = -x[base + half + d];
            rot[half + d] = x[base + d];
        }
        for d in 0..hd {
            x[base + d] = x[base + d] * cos[d] + rot[d] * sin[d];
        }
    }
}

struct DecoderLayer {
    input_ln: Vec<f32>,
    q: PackedWeight,  // [heads*hd, h]
    k: PackedWeight,  // [kv*hd, h]
    v: PackedWeight,  // [kv*hd, h]
    o: PackedWeight,  // [h, heads*hd]
    q_norm: Vec<f32>, // [hd]
    k_norm: Vec<f32>, // [hd]
    post_ln: Vec<f32>,
    gate: PackedWeight,
    up: PackedWeight,
    down: PackedWeight,
}

/// GEMV through the prepacked panels (`cpu_gemm::gemm_packed`, m=1) — ~3× the naive strict-order
/// dot per the moshi_lm measurements; the raw row-major copy is dropped at load
/// (prepack-and-drop, landmine §5.11).
fn pw_mul(w: &PackedWeight, x: &[f32]) -> Vec<f32> {
    debug_assert_eq!(x.len(), w.k());
    let mut out = vec![0f32; w.n()];
    crate::cpu_gemm::gemm_packed(&mut out, x, w, 1, None);
    out
}

fn pack(v: Vec<f32>, n: usize, k: usize) -> PackedWeight {
    PackedWeight::new(&v, n, k)
}

/// Dims of one Qwen3-style decoder stack (talker or code predictor — same math, different sizes).
#[derive(Debug, Clone, Copy)]
struct StackDims {
    h: usize,
    layers: usize,
    heads: usize,
    kv_heads: usize,
    hd: usize,
    inter: usize,
    eps: f32,
    theta: f32,
}

/// Per-layer KV cache: `k[layer]`/`v[layer]` are `[len × kv_width]` row-major.
pub struct TalkerState {
    k: Vec<Vec<f32>>,
    v: Vec<Vec<f32>>,
    len: usize,
}

impl TalkerState {
    pub fn len(&self) -> usize {
        self.len
    }
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }
}

/// The generic Qwen3-style stack forward (shared by Talker now; CodePredictor reuses it at P3).
struct Stack {
    dims: StackDims,
    layers: Vec<DecoderLayer>,
    final_norm: Vec<f32>,
    axes: Option<Vec<u8>>, // Some(mrope axis map) or None → 1-D rope
}

impl Stack {
    fn load(
        st: &chatterbox_st::St,
        prefix: &str,
        dims: StackDims,
        mrope: Option<[usize; 3]>,
    ) -> Result<Self> {
        let (h, qw, kvw) = (dims.h, dims.heads * dims.hd, dims.kv_heads * dims.hd);
        let mut layers = Vec::with_capacity(dims.layers);
        for l in 0..dims.layers {
            let p = format!("{prefix}.layers.{l}");
            layers.push(DecoderLayer {
                input_ln: st.f32(&format!("{p}.input_layernorm.weight"))?,
                q: pack(
                    st.mat(&format!("{p}.self_attn.q_proj.weight"), qw, h)?,
                    qw,
                    h,
                ),
                k: pack(
                    st.mat(&format!("{p}.self_attn.k_proj.weight"), kvw, h)?,
                    kvw,
                    h,
                ),
                v: pack(
                    st.mat(&format!("{p}.self_attn.v_proj.weight"), kvw, h)?,
                    kvw,
                    h,
                ),
                o: pack(
                    st.mat(&format!("{p}.self_attn.o_proj.weight"), h, qw)?,
                    h,
                    qw,
                ),
                q_norm: st.f32(&format!("{p}.self_attn.q_norm.weight"))?,
                k_norm: st.f32(&format!("{p}.self_attn.k_norm.weight"))?,
                post_ln: st.f32(&format!("{p}.post_attention_layernorm.weight"))?,
                gate: pack(
                    st.mat(&format!("{p}.mlp.gate_proj.weight"), dims.inter, h)?,
                    dims.inter,
                    h,
                ),
                up: pack(
                    st.mat(&format!("{p}.mlp.up_proj.weight"), dims.inter, h)?,
                    dims.inter,
                    h,
                ),
                down: pack(
                    st.mat(&format!("{p}.mlp.down_proj.weight"), h, dims.inter)?,
                    h,
                    dims.inter,
                ),
            });
        }
        Ok(Self {
            dims,
            layers,
            final_norm: st.f32(&format!("{prefix}.norm.weight"))?,
            axes: mrope.map(|s| mrope_axis_map(s, dims.hd / 2)),
        })
    }

    fn rope_at(&self, pos: u32) -> (Vec<f32>, Vec<f32>) {
        match &self.axes {
            // Uniform positions by construction (reference `get_rope_index` is a plain cumsum
            // over all three channels) — the runtime assert of landmine §5.1 lives in the fact
            // that only [p,p,p] is constructible here.
            Some(a) => mrope_cs(self.dims.hd, self.dims.theta, a, [pos, pos, pos]),
            None => rope_cs_1d(self.dims.hd, self.dims.theta, pos),
        }
    }

    fn new_state(&self) -> TalkerState {
        TalkerState {
            k: vec![Vec::new(); self.dims.layers],
            v: vec![Vec::new(); self.dims.layers],
            len: 0,
        }
    }

    /// Append ONE position (embed `[h]`, absolute `pos`) to the cache; return its final-norm
    /// hidden `[h]`. `prefill` = fold over this; the no-cache reference must match bit-exactly.
    fn step(&self, state: &mut TalkerState, embed: &[f32], pos: u32) -> Vec<f32> {
        let d = self.dims;
        let (h, hd, qw, kvw) = (d.h, d.hd, d.heads * d.hd, d.kv_heads * d.hd);
        let group = d.heads / d.kv_heads;
        let scale = 1.0 / (hd as f32).sqrt();
        let (cos, sin) = self.rope_at(pos);
        let mut x = embed.to_vec();
        for (l, layer) in self.layers.iter().enumerate() {
            let xn = rms_norm(&x, &layer.input_ln, d.eps);
            let mut q = pw_mul(&layer.q, &xn);
            let mut kx = pw_mul(&layer.k, &xn);
            let vx = pw_mul(&layer.v, &xn);
            // Per-head qk RMS-norm, then rope.
            for head in 0..d.heads {
                let s = &mut q[head * hd..(head + 1) * hd];
                s.copy_from_slice(&rms_norm(s, &layer.q_norm, d.eps));
            }
            for head in 0..d.kv_heads {
                let s = &mut kx[head * hd..(head + 1) * hd];
                s.copy_from_slice(&rms_norm(s, &layer.k_norm, d.eps));
            }
            apply_rope(&mut q, &cos, &sin, d.heads, hd);
            apply_rope(&mut kx, &cos, &sin, d.kv_heads, hd);
            state.k[l].extend_from_slice(&kx);
            state.v[l].extend_from_slice(&vx);
            let t_len = state.k[l].len() / kvw;
            // GQA attention of this single query row over the cache.
            let mut attn = vec![0f32; qw];
            for head in 0..d.heads {
                let kvh = head / group;
                let qh = &q[head * hd..(head + 1) * hd];
                let mut scores = vec![0f32; t_len];
                for (t, sc) in scores.iter_mut().enumerate() {
                    let kh = &state.k[l][t * kvw + kvh * hd..t * kvw + kvh * hd + hd];
                    *sc = qh.iter().zip(kh).map(|(a, b)| a * b).sum::<f32>() * scale;
                }
                softmax_inplace(&mut scores);
                let out = &mut attn[head * hd..(head + 1) * hd];
                for (t, &w) in scores.iter().enumerate() {
                    let vh = &state.v[l][t * kvw + kvh * hd..t * kvw + kvh * hd + hd];
                    for (dd, ov) in out.iter_mut().enumerate() {
                        *ov += w * vh[dd];
                    }
                }
            }
            let o = pw_mul(&layer.o, &attn);
            for j in 0..h {
                x[j] += o[j];
            }
            let xn = rms_norm(&x, &layer.post_ln, d.eps);
            let g = pw_mul(&layer.gate, &xn);
            let u = pw_mul(&layer.up, &xn);
            let act: Vec<f32> = g.iter().zip(&u).map(|(gv, uv)| silu(*gv) * uv).collect();
            let dn = pw_mul(&layer.down, &act);
            for j in 0..h {
                x[j] += dn[j];
            }
        }
        state.len += 1;
        rms_norm(&x, &self.final_norm, d.eps)
    }

    /// Whole-sequence causal forward WITHOUT a cache — the reference the cached path is gated
    /// against. Returns final-norm hiddens `[seq, h]`. Positions are `pos0..pos0+seq`.
    fn forward_hidden(&self, embeds: &[f32], pos0: u32) -> Vec<f32> {
        let d = self.dims;
        let (h, hd, qw, kvw) = (d.h, d.hd, d.heads * d.hd, d.kv_heads * d.hd);
        let group = d.heads / d.kv_heads;
        let seq = embeds.len() / h;
        let scale = 1.0 / (hd as f32).sqrt();
        let tables: Vec<(Vec<f32>, Vec<f32>)> =
            (0..seq).map(|s| self.rope_at(pos0 + s as u32)).collect();
        let mut x = embeds.to_vec();
        for layer in &self.layers {
            let mut q = vec![0f32; seq * qw];
            let mut k = vec![0f32; seq * kvw];
            let mut v = vec![0f32; seq * kvw];
            for s in 0..seq {
                let xn = rms_norm(&x[s * h..(s + 1) * h], &layer.input_ln, d.eps);
                let qs = pw_mul(&layer.q, &xn);
                let ks = pw_mul(&layer.k, &xn);
                let vs = pw_mul(&layer.v, &xn);
                q[s * qw..(s + 1) * qw].copy_from_slice(&qs);
                k[s * kvw..(s + 1) * kvw].copy_from_slice(&ks);
                v[s * kvw..(s + 1) * kvw].copy_from_slice(&vs);
                let (cos, sin) = &tables[s];
                for head in 0..d.heads {
                    let sl = &mut q[s * qw + head * hd..s * qw + (head + 1) * hd];
                    sl.copy_from_slice(&rms_norm(sl, &layer.q_norm, d.eps));
                }
                for head in 0..d.kv_heads {
                    let sl = &mut k[s * kvw + head * hd..s * kvw + (head + 1) * hd];
                    sl.copy_from_slice(&rms_norm(sl, &layer.k_norm, d.eps));
                }
                apply_rope(&mut q[s * qw..(s + 1) * qw], cos, sin, d.heads, hd);
                apply_rope(&mut k[s * kvw..(s + 1) * kvw], cos, sin, d.kv_heads, hd);
            }
            let mut attn = vec![0f32; seq * qw];
            for head in 0..d.heads {
                let kvh = head / group;
                for s in 0..seq {
                    let qh = &q[s * qw + head * hd..s * qw + (head + 1) * hd];
                    let mut scores = vec![0f32; s + 1];
                    for (t, sc) in scores.iter_mut().enumerate() {
                        let kh = &k[t * kvw + kvh * hd..t * kvw + kvh * hd + hd];
                        *sc = qh.iter().zip(kh).map(|(a, b)| a * b).sum::<f32>() * scale;
                    }
                    softmax_inplace(&mut scores);
                    let out = &mut attn[s * qw + head * hd..s * qw + (head + 1) * hd];
                    for (t, &w) in scores.iter().enumerate() {
                        let vh = &v[t * kvw + kvh * hd..t * kvw + kvh * hd + hd];
                        for (dd, ov) in out.iter_mut().enumerate() {
                            *ov += w * vh[dd];
                        }
                    }
                }
            }
            for s in 0..seq {
                let o = pw_mul(&layer.o, &attn[s * qw..(s + 1) * qw]);
                for j in 0..h {
                    x[s * h + j] += o[j];
                }
                let xn = rms_norm(&x[s * h..(s + 1) * h], &layer.post_ln, d.eps);
                let g = pw_mul(&layer.gate, &xn);
                let u = pw_mul(&layer.up, &xn);
                let act: Vec<f32> = g.iter().zip(&u).map(|(gv, uv)| silu(*gv) * uv).collect();
                let dn = pw_mul(&layer.down, &act);
                for j in 0..h {
                    x[s * h + j] += dn[j];
                }
            }
        }
        (0..seq)
            .flat_map(|s| rms_norm(&x[s * h..(s + 1) * h], &self.final_norm, d.eps))
            .collect()
    }
}

/// Thin named wrapper over the chatterbox safetensors reader (shared manifest tooling) adding the
/// shape-asserting `mat` this module needs.
pub(crate) mod chatterbox_st {
    use anyhow::Result;
    use std::path::Path;

    pub struct St(crate::chatterbox::StReader);

    impl St {
        pub fn open(path: &Path) -> Result<Self> {
            Ok(Self(crate::chatterbox::StReader::open(path)?))
        }
        pub fn f32(&self, name: &str) -> Result<Vec<f32>> {
            self.0.f32(name)
        }
        pub fn shape(&self, name: &str) -> Result<&[usize]> {
            self.0.shape(name)
        }
        pub fn mat(&self, name: &str, rows: usize, cols: usize) -> Result<Vec<f32>> {
            let s = self.0.shape(name)?;
            anyhow::ensure!(s == [rows, cols], "{name}: shape {s:?} != [{rows}, {cols}]");
            self.0.f32(name)
        }
    }
}

/// The Talker: dual-stream embedding tables + the Qwen3 stack + the untied codec head.
pub struct Talker {
    stack: Stack,
    /// [151936, th] — ids are always < 151676 (base+added); the rest is padding. The text table
    /// is `text_hidden_size` wide (2048 on BOTH sizes — the 0.6B's ResizeMLP genuinely resizes
    /// 2048→1024; on the 1.7B it happens to be square).
    text_embedding: Vec<f32>,
    /// [3072, h] — codes 0..2047 + specials 2048+.
    codec_embedding: Vec<f32>,
    text_fc1_w: PackedWeight, // [th, th]
    text_fc1_b: Vec<f32>,
    text_fc2_w: PackedWeight, // [h, th]
    text_fc2_b: Vec<f32>,
    codec_head: PackedWeight, // [3072, h]
    h: usize,
    th: usize,
}

impl Talker {
    /// Load the talker stack + embeds/head from `<dir>/model.safetensors` (bf16 → f32).
    pub fn load(dir: &Path, config: &Qwen3TtsConfig) -> Result<Self> {
        let t = &config.top.talker_config;
        let st = chatterbox_st::St::open(&dir.join(cfg::FILE_MODEL))?;
        let dims = StackDims {
            h: t.hidden_size,
            layers: t.num_hidden_layers,
            heads: t.num_attention_heads,
            kv_heads: t.num_key_value_heads,
            hd: t.head_dim,
            inter: t.intermediate_size,
            eps: t.rms_norm_eps,
            theta: t.rope_theta,
        };
        let stack = Stack::load(&st, "talker.model", dims, Some(t.mrope_section))?;
        let h = t.hidden_size;
        let th = t.text_hidden_size;
        Ok(Self {
            stack,
            text_embedding: st.mat("talker.model.text_embedding.weight", t.text_vocab_size, th)?,
            codec_embedding: st.mat("talker.model.codec_embedding.weight", t.vocab_size, h)?,
            text_fc1_w: pack(
                st.mat("talker.text_projection.linear_fc1.weight", th, th)?,
                th,
                th,
            ),
            text_fc1_b: st.f32("talker.text_projection.linear_fc1.bias")?,
            text_fc2_w: pack(
                st.mat("talker.text_projection.linear_fc2.weight", h, th)?,
                h,
                th,
            ),
            text_fc2_b: st.f32("talker.text_projection.linear_fc2.bias")?,
            codec_head: pack(
                st.mat("talker.codec_head.weight", t.vocab_size, h)?,
                t.vocab_size,
                h,
            ),
            h,
            th,
        })
    }

    /// `text_projection(text_embedding[id])` — the ResizeMLP `fc2(act(fc1(x)))` with the
    /// checkpoint's silu activation (`talker_config.hidden_act`). Genuinely a resize on the 0.6B
    /// (2048→1024); square on the 1.7B.
    pub fn embed_text(&self, id: u32) -> Vec<f32> {
        let e = &self.text_embedding[id as usize * self.th..(id as usize + 1) * self.th];
        self.project_text(e)
    }

    /// The same ResizeMLP over an arbitrary text-hidden row (used by P6 for summed streams).
    pub fn project_text(&self, e: &[f32]) -> Vec<f32> {
        let mut a = vec![0f32; self.th];
        crate::cpu_gemm::gemm_packed(&mut a, e, &self.text_fc1_w, 1, Some(&self.text_fc1_b));
        for v in a.iter_mut() {
            *v = silu(*v);
        }
        let mut o = vec![0f32; self.h];
        crate::cpu_gemm::gemm_packed(&mut o, &a, &self.text_fc2_w, 1, Some(&self.text_fc2_b));
        o
    }

    pub fn embed_codec(&self, id: u32) -> &[f32] {
        &self.codec_embedding[id as usize * self.h..(id as usize + 1) * self.h]
    }

    pub fn new_state(&self) -> TalkerState {
        self.stack.new_state()
    }

    /// One cached step; returns the final-norm hidden `[h]` of the appended position.
    pub fn step(&self, state: &mut TalkerState, embed: &[f32], pos: u32) -> Vec<f32> {
        self.stack.step(state, embed, pos)
    }

    /// Cached prefill of a whole `[seq, h]` embed block; returns the LAST position's hidden.
    pub fn prefill(&self, state: &mut TalkerState, embeds: &[f32]) -> Vec<f32> {
        let seq = embeds.len() / self.h;
        let mut last = Vec::new();
        for s in 0..seq {
            let pos = state.len as u32;
            last = self
                .stack
                .step(state, &embeds[s * self.h..(s + 1) * self.h], pos);
        }
        last
    }

    /// No-cache whole-sequence reference forward (final-norm hiddens `[seq, h]`).
    pub fn forward_hidden(&self, embeds: &[f32], pos0: u32) -> Vec<f32> {
        self.stack.forward_hidden(embeds, pos0)
    }

    /// Codec-head logits `[3072]` over one hidden row.
    pub fn codec_logits(&self, hidden: &[f32]) -> Vec<f32> {
        pw_mul(&self.codec_head, hidden)
    }

    pub fn hidden_size(&self) -> usize {
        self.h
    }
}

/// Exposed for the P2 identity gate.
pub fn mrope_uniform_tables(
    hd: usize,
    theta: f32,
    section: [usize; 3],
    pos: u32,
) -> (Vec<f32>, Vec<f32>) {
    let axes = mrope_axis_map(section, hd / 2);
    mrope_cs(hd, theta, &axes, [pos, pos, pos])
}

/// Exposed for the P2 identity gate.
pub fn rope_1d_tables(hd: usize, theta: f32, pos: u32) -> (Vec<f32>, Vec<f32>) {
    rope_cs_1d(hd, theta, pos)
}

// ---------------------------------------------------------------------------------------------
// P3 — CodePredictor + the AR frame loop.
//
// Wiring transcribed VERBATIM from the reference `modeling_qwen3_tts.py` (2026-07-20):
//   - per frame the talker samples ONE semantic code c0 from `codec_head(last_hidden)`
//     (HF repetition penalty 1.05 over the talker's codec-side id stream: prompt ids ∪ previous
//     semantic codes — subtalker sampling has NO rep-pen);
//   - the CodePredictor is a fresh `generate` per frame (fresh DynamicCache — KV RESET, gated by
//     `p3_code_predictor_kv_reset_invariance`): prefill = [past_hidden, codec_embedding[c0]]
//     (both TALKER-width), every input passes `small_to_mtp_projection` into the 1024 core;
//     15 steps: head[k] predicts code group k+1, its code re-enters via cp table[k] (talker-width)
//     at the next position — `lm_head[i-1] reads position i`, `get_input_embeddings()[gs-1]`;
//   - next talker input = `cat([embed(c0)] + [cp_table[i](code[i+1])]).sum(1)` — the SUM of all
//     16 talker-width code embeddings — plus `trailing_text_hidden[step]` while text remains,
//     else `tts_pad_embed` (landmine §5.7: the off-by-one here degrades quietly);
//   - stop: c0 == codec_eos (2150); the CodePredictor is NOT run on the stop frame.
// ---------------------------------------------------------------------------------------------

/// Sampling knobs for one stream (talker or subtalker), mapped 1:1 onto `crate::sampling`.
/// The draw RNG is the engine's seeded xorshift64, NOT torch Philox — documented deviation
/// (HANDOFF_qwen3tts.md §2): there is no oracle to match, seeded determinism is what the gates
/// need, and the truncation chain (temp → top-k → top-p) is order-identical to HF's.
#[derive(Debug, Clone, Copy)]
pub struct SamplingCfg {
    pub temperature: f32,
    pub top_k: usize,
    pub top_p: f32,
    /// HF multiplicative repetition penalty over prompt ∪ emitted; 1.0 = off.
    pub repetition_penalty: f32,
}

impl SamplingCfg {
    /// The checkpoint's own talker/subtalker defaults.
    pub fn from_generation(g: &GenerationConfig) -> (Self, Self) {
        (
            Self {
                temperature: g.temperature,
                top_k: g.top_k,
                top_p: g.top_p,
                repetition_penalty: g.repetition_penalty,
            },
            Self {
                temperature: g.subtalker_temperature,
                top_k: g.subtalker_top_k,
                top_p: g.subtalker_top_p,
                repetition_penalty: 1.0,
            },
        )
    }
}

/// The 5-layer MTP decoder that turns one talker hidden + one semantic code into the 15 acoustic
/// codes of a frame. Its 15 embedding tables are TALKER-width (they also serve the frame-sum
/// feedback into the talker); its core runs at 1024 behind `small_to_mtp_projection`.
pub struct CodePredictor {
    stack: Stack,
    /// 15 × [2048, talker_h] — table `i` embeds code GROUP `i+1`.
    tables: Vec<Vec<f32>>,
    /// 15 × [2048, cp_h] — head `k` predicts code group `k+1`.
    heads: Vec<PackedWeight>,
    /// `small_to_mtp_projection` [cp_h, talker_h] + bias — ABSENT when cp_h == talker_h (the
    /// 0.6B: the reference only instantiates it for a genuine resize; identity otherwise).
    proj: Option<(PackedWeight, Vec<f32>)>,
    talker_h: usize,
    cp_h: usize,
}

impl CodePredictor {
    pub fn load(dir: &Path, config: &Qwen3TtsConfig) -> Result<Self> {
        let t = &config.top.talker_config;
        let p = &t.code_predictor_config;
        let st = chatterbox_st::St::open(&dir.join(cfg::FILE_MODEL))?;
        let dims = StackDims {
            h: p.hidden_size,
            layers: p.num_hidden_layers,
            heads: p.num_attention_heads,
            kv_heads: p.num_key_value_heads,
            hd: p.head_dim,
            inter: p.intermediate_size,
            eps: p.rms_norm_eps,
            theta: p.rope_theta,
        };
        let stack = Stack::load(&st, "talker.code_predictor.model", dims, None)?;
        let n = p.num_code_groups - 1;
        let mut tables = Vec::with_capacity(n);
        let mut heads = Vec::with_capacity(n);
        for g in 0..n {
            tables.push(st.mat(
                &format!("talker.code_predictor.model.codec_embedding.{g}.weight"),
                p.vocab_size,
                t.hidden_size,
            )?);
            heads.push(pack(
                st.mat(
                    &format!("talker.code_predictor.lm_head.{g}.weight"),
                    p.vocab_size,
                    p.hidden_size,
                )?,
                p.vocab_size,
                p.hidden_size,
            ));
        }
        let proj = if p.hidden_size == t.hidden_size {
            anyhow::ensure!(
                st.shape("talker.code_predictor.small_to_mtp_projection.weight")
                    .is_err(),
                "cp_h == talker_h but a small_to_mtp_projection exists — layout drift"
            );
            None
        } else {
            Some((
                pack(
                    st.mat(
                        "talker.code_predictor.small_to_mtp_projection.weight",
                        p.hidden_size,
                        t.hidden_size,
                    )?,
                    p.hidden_size,
                    t.hidden_size,
                ),
                st.f32("talker.code_predictor.small_to_mtp_projection.bias")?,
            ))
        };
        Ok(Self {
            stack,
            tables,
            heads,
            proj,
            talker_h: t.hidden_size,
            cp_h: p.hidden_size,
        })
    }

    fn project(&self, x: &[f32]) -> Vec<f32> {
        match &self.proj {
            None => x.to_vec(),
            Some((w, b)) => {
                let mut o = vec![0f32; self.cp_h];
                crate::cpu_gemm::gemm_packed(&mut o, x, w, 1, Some(b));
                o
            }
        }
    }

    /// Talker-width embedding row of code GROUP `group` (1..=15) — table `group-1`.
    pub fn embed_group(&self, group: usize, code: u32) -> &[f32] {
        debug_assert!((1..cfg::NUM_CODE_GROUPS).contains(&group));
        let t = &self.tables[group - 1];
        &t[code as usize * self.talker_h..(code as usize + 1) * self.talker_h]
    }

    /// One frame: fresh KV, prefill `[proj(past_hidden), proj(embed_c0)]`, then 15 sampled steps.
    /// Returns the acoustic codes for groups 1..=15 (each < 2048).
    pub fn predict_frame(
        &self,
        past_hidden: &[f32],
        sem_embed: &[f32],
        sampling: &SamplingCfg,
        rng: &mut u64,
    ) -> Vec<u32> {
        let mut state = self.stack.new_state();
        self.stack.step(&mut state, &self.project(past_hidden), 0);
        let mut h = self.stack.step(&mut state, &self.project(sem_embed), 1);
        let mut codes = Vec::with_capacity(cfg::NUM_CODE_GROUPS - 1);
        for k in 0..cfg::NUM_CODE_GROUPS - 1 {
            let logits = pw_mul(&self.heads[k], &h);
            let code = crate::sampling::sample(
                &logits,
                sampling.temperature,
                Some(sampling.top_k),
                sampling.top_p,
                None,
                rng,
            );
            codes.push(code);
            if k + 1 < cfg::NUM_CODE_GROUPS - 1 {
                let e = self.embed_group(k + 1, code).to_vec();
                h = self
                    .stack
                    .step(&mut state, &self.project(&e), (k + 2) as u32);
            }
        }
        codes
    }
}

/// One generated frame of the AR loop, `StepTrace`-style (moshi_lm precedent): enough to localize
/// a divergence per frame instead of smearing it across the clip.
#[derive(Debug, Clone, PartialEq)]
pub struct FrameTrace {
    /// Semantic code (group 0) — may be `codec_eos`, in which case `acoustic` is empty and this
    /// is the final frame.
    pub semantic: u32,
    /// Groups 1..=15, each < 2048. Empty on the stop frame.
    pub acoustic: Vec<u32>,
}

/// The talker AR frame loop over an already-assembled prompt (P6 owns prompt assembly; the loop
/// takes the prefill embeds + the projected trailing-text stream + the projected tts_pad row).
///
/// `prompt_codec_ids` seeds the repetition-penalty id set. NOTE (transcribed from `generate()`):
/// the reference drives HF `generate` with `inputs_embeds` ONLY, so its penalty stream contains
/// GENERATED semantic codes exclusively — faithful callers pass `&[]` here. The reference also
/// suppresses every special id (`vocab−1024..vocab` except `codec_eos`) at every step, and
/// forbids EOS for the first two frames (`min_new_tokens: 2`) — both applied below.
#[allow(clippy::too_many_arguments)]
pub fn generate_frames(
    talker: &Talker,
    cp: &CodePredictor,
    prompt_embeds: &[f32],
    prompt_codec_ids: &[u32],
    trailing_text: &[Vec<f32>],
    tts_pad_embed: &[f32],
    talker_sampling: &SamplingCfg,
    subtalker_sampling: &SamplingCfg,
    seed: u64,
    max_frames: usize,
) -> Vec<FrameTrace> {
    let mut frames = Vec::new();
    generate_frames_cb(
        talker,
        cp,
        prompt_embeds,
        prompt_codec_ids,
        trailing_text,
        tts_pad_embed,
        talker_sampling,
        subtalker_sampling,
        seed,
        max_frames,
        &mut |f| {
            frames.push(f);
            true
        },
    );
    frames
}

/// [`generate_frames`] with a per-frame callback (return `false` to abort early) — the streaming
/// service path decodes and ships each 80 ms frame as it is sampled.
#[allow(clippy::too_many_arguments)]
pub fn generate_frames_cb(
    talker: &Talker,
    cp: &CodePredictor,
    prompt_embeds: &[f32],
    prompt_codec_ids: &[u32],
    trailing_text: &[Vec<f32>],
    tts_pad_embed: &[f32],
    talker_sampling: &SamplingCfg,
    subtalker_sampling: &SamplingCfg,
    seed: u64,
    max_frames: usize,
    on_frame: &mut dyn FnMut(FrameTrace) -> bool,
) {
    const MIN_NEW_FRAMES: usize = 2;
    let h = talker.hidden_size();
    let mut rng = if seed == 0 { 0x9E3779B97F4A7C15 } else { seed };
    let mut state = talker.new_state();
    let mut last_hidden = talker.prefill(&mut state, prompt_embeds);
    let mut emitted: Vec<u32> = Vec::new();
    for step in 0..max_frames {
        let mut logits = talker.codec_logits(&last_hidden);
        crate::sampling::apply_penalties(
            &mut logits,
            prompt_codec_ids,
            &emitted,
            0.0,
            0.0,
            talker_sampling.repetition_penalty,
        );
        // suppress_tokens: specials 2048..3072 can never be sampled — except EOS, which is also
        // forbidden until MIN_NEW_FRAMES frames exist.
        for id in cfg::CP_VOCAB..cfg::TALKER_CODEC_VOCAB {
            if id as u32 != cfg::CODEC_EOS || emitted.len() < MIN_NEW_FRAMES {
                logits[id] = f32::NEG_INFINITY;
            }
        }
        let c0 = crate::sampling::sample(
            &logits,
            talker_sampling.temperature,
            Some(talker_sampling.top_k),
            talker_sampling.top_p,
            None,
            &mut rng,
        );
        emitted.push(c0);
        if c0 == cfg::CODEC_EOS {
            on_frame(FrameTrace {
                semantic: c0,
                acoustic: Vec::new(),
            });
            break;
        }
        let sem_embed = talker.embed_codec(c0).to_vec();
        let acoustic = cp.predict_frame(&last_hidden, &sem_embed, subtalker_sampling, &mut rng);
        // Next talker input: Σ of all 16 talker-width code embeds + trailing text (or tts_pad).
        let mut next = sem_embed.clone();
        for (i, &code) in acoustic.iter().enumerate() {
            let e = cp.embed_group(i + 1, code);
            for (n, v) in next.iter_mut().zip(e) {
                *n += *v;
            }
        }
        let text_row: &[f32] = trailing_text
            .get(step)
            .map(|v| v.as_slice())
            .unwrap_or(tts_pad_embed);
        for (n, v) in next.iter_mut().zip(text_row) {
            *n += *v;
        }
        debug_assert_eq!(next.len(), h);
        if !on_frame(FrameTrace {
            semantic: c0,
            acoustic,
        }) {
            break;
        }
        let pos = state.len() as u32;
        last_hidden = talker.step(&mut state, &next, pos);
    }
}

// ---------------------------------------------------------------------------------------------
// P6b — prompt assembly, transcribed VERBATIM from `generate()` in the reference
// `modeling_qwen3_tts.py` (lines ~2124–2234, read directly 2026-07-20):
//
//   text template  : "<|im_start|>assistant\n{TEXT}<|im_end|>\n<|im_start|>assistant\n"
//                    → ids[..3] role, ids[3..len−5] text, 5 trailing template tokens;
//   ref template   : "<|im_start|>assistant\n{REF}<|im_end|>\n" → ids[3..len−2] ref text.
//
//   prefill rows (streaming mode, the reference default):
//     [0..3)   proj(text_embed(role))                       — text stream only
//     next L−1 [tts_pad×(L−2), tts_bos] + codec[..L−1]      — codec = prefill specials
//              (auto: [nothink, think_bos, think_eos]; lang: [think, think_bos, lang_id,
//              think_eos]) ⧺ optional RAW x-vector row ⧺ [codec_pad, codec_bos]
//     non-ICL: 1 row proj(text_embed(ids[3])) + codec[L−1]  — first text token over codec_bos;
//              trailing = proj(text_embed(ids[4..len−5])) ⧺ tts_eos
//     ICL:     text_rows = proj(ref_ids[3..len−2] ⧺ ids[3..len−5]) ⧺ tts_eos   (T1)
//              codec_rows = [codec_bos] ⧺ frame-sum(ref_codes[t]) per ref frame (T2)
//              T1 > T2 → rows += text[..T2] + codec, trailing = text[T2..]
//              else    → rows += (text ⧺ tts_pad×(T2−T1)) + codec, trailing = ∅ (pads follow)
// ---------------------------------------------------------------------------------------------

/// An assembled prefill for [`generate_frames`].
pub struct TtsPrompt {
    /// `[seq, h]` prefill embeddings.
    pub embeds: Vec<f32>,
    /// Projected trailing-text rows (ends with the tts_eos row; empty when the ICL pad case
    /// consumed the whole text inside the prefill).
    pub trailing_text: Vec<Vec<f32>>,
    /// The projected tts_pad row `generate_frames` falls back to after the text runs out.
    pub tts_pad: Vec<f32>,
}

/// The exact tokenized template of the reference pipeline.
pub fn wrap_text_template(text: &str) -> String {
    format!("<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n")
}

fn wrap_ref_template(text: &str) -> String {
    format!("<|im_start|>assistant\n{text}<|im_end|>\n")
}

/// Base-checkpoint prompt: zero-shot clone via `x_vector` (raw speaker-encoder output), optional
/// ICL `(ref_text, ref_codes)` (codec-encoder output of the reference clip). `language: None` =
/// auto-detect (the nothink prefill).
#[cfg(feature = "cli")]
#[allow(clippy::too_many_arguments)]
pub fn build_clone_prompt(
    talker: &Talker,
    cp: &CodePredictor,
    tk: &TextTokenizer,
    config: &Qwen3TtsConfig,
    text: &str,
    language: Option<&str>,
    x_vector: Option<&[f32]>,
    icl: Option<(&str, &[[u32; cfg::NUM_CODE_GROUPS]])>,
) -> Result<TtsPrompt> {
    let lang_id = resolve_language(config, language)?;
    let speaker_row = x_vector.map(|v| v.to_vec());
    build_prompt_inner(
        talker,
        cp,
        tk,
        config,
        text,
        lang_id,
        speaker_row,
        icl,
        None,
    )
}

/// CustomVoice: preset `speaker` (spk_id → codec-embedding row) + optional style `instruct`
/// (template `<|im_start|>user\n{instruct}<|im_end|>\n`, prepended as text-only rows). The
/// dialect override is the reference's: chinese/auto + a dialect speaker → dialect language id.
#[cfg(feature = "cli")]
pub fn build_preset_prompt(
    talker: &Talker,
    cp: &CodePredictor,
    tk: &TextTokenizer,
    config: &Qwen3TtsConfig,
    text: &str,
    language: Option<&str>,
    speaker: &str,
    instruct: Option<&str>,
) -> Result<TtsPrompt> {
    let t = &config.top.talker_config;
    let key = speaker.to_lowercase();
    let spk_id = *t
        .spk_id
        .get(&key)
        .with_context(|| format!("speaker {speaker:?} not implemented"))?;
    let mut lang_id = resolve_language(config, language)?;
    let auto_or_zh = matches!(
        language.map(|l| l.to_lowercase()).as_deref(),
        None | Some("auto") | Some("chinese")
    );
    if auto_or_zh && let Some(Some(dialect)) = t.spk_is_dialect.get(&key) {
        lang_id = Some(
            *t.codec_language_id
                .get(dialect)
                .with_context(|| format!("dialect {dialect:?} missing from codec_language_id"))?,
        );
    }
    let speaker_row = Some(talker.embed_codec(spk_id).to_vec());
    build_prompt_inner(
        talker,
        cp,
        tk,
        config,
        text,
        lang_id,
        speaker_row,
        None,
        instruct,
    )
}

/// VoiceDesign: a natural-language voice description as `instruct`, no speaker row.
#[cfg(feature = "cli")]
pub fn build_design_prompt(
    talker: &Talker,
    cp: &CodePredictor,
    tk: &TextTokenizer,
    config: &Qwen3TtsConfig,
    text: &str,
    language: Option<&str>,
    instruct: &str,
) -> Result<TtsPrompt> {
    let lang_id = resolve_language(config, language)?;
    build_prompt_inner(
        talker,
        cp,
        tk,
        config,
        text,
        lang_id,
        None,
        None,
        Some(instruct),
    )
}

fn resolve_language(config: &Qwen3TtsConfig, language: Option<&str>) -> Result<Option<u32>> {
    match language {
        None => Ok(None),
        Some(l) if l.eq_ignore_ascii_case("auto") => Ok(None),
        Some(l) => Ok(Some(
            *config
                .top
                .talker_config
                .codec_language_id
                .get(&l.to_lowercase())
                .with_context(|| format!("language {l:?} not implemented"))?,
        )),
    }
}

#[cfg(feature = "cli")]
#[allow(clippy::too_many_arguments)]
fn build_prompt_inner(
    talker: &Talker,
    cp: &CodePredictor,
    tk: &TextTokenizer,
    config: &Qwen3TtsConfig,
    text: &str,
    lang_id: Option<u32>,
    speaker_row: Option<Vec<f32>>,
    icl: Option<(&str, &[[u32; cfg::NUM_CODE_GROUPS]])>,
    instruct: Option<&str>,
) -> Result<TtsPrompt> {
    let h = talker.hidden_size();
    let _ = &config.top.talker_config;
    let ids = tk.encode(&wrap_text_template(text))?;
    anyhow::ensure!(
        ids.len() >= 9,
        "text tokenizes too short: {} ids",
        ids.len()
    );
    let text_ids = &ids[3..ids.len() - 5];

    let tts_bos = talker.embed_text(cfg::TTS_BOS);
    let tts_eos = talker.embed_text(cfg::TTS_EOS);
    let tts_pad = talker.embed_text(cfg::TTS_PAD);

    // Codec stream: prefill specials ⧺ speaker row (raw x-vector OR spk_id codec row) ⧺ [pad, bos].
    let mut codec_rows: Vec<Vec<f32>> = Vec::new();
    match lang_id {
        None => {
            for id in [
                cfg::CODEC_NOTHINK,
                cfg::CODEC_THINK_BOS,
                cfg::CODEC_THINK_EOS,
            ] {
                codec_rows.push(talker.embed_codec(id).to_vec());
            }
        }
        Some(lang_id) => {
            for id in [
                cfg::CODEC_THINK,
                cfg::CODEC_THINK_BOS,
                lang_id,
                cfg::CODEC_THINK_EOS,
            ] {
                codec_rows.push(talker.embed_codec(id).to_vec());
            }
        }
    }
    if let Some(row) = speaker_row {
        anyhow::ensure!(row.len() == h, "speaker row dim {} ≠ hidden {h}", row.len());
        codec_rows.push(row);
    }
    codec_rows.push(talker.embed_codec(cfg::CODEC_PAD).to_vec());
    codec_rows.push(talker.embed_codec(cfg::CODEC_BOS).to_vec());
    let l = codec_rows.len();

    let mut embeds: Vec<f32> = Vec::new();
    // Instruct rows come FIRST — text-only, whole template, no slicing (reference: generate()
    // appends proj(text_embed(instruct_id)) before the tts prompt).
    if let Some(instr) = instruct
        && !instr.trim().is_empty()
    {
        for id in tk.encode(&format!("<|im_start|>user\n{instr}<|im_end|>\n"))? {
            embeds.extend(talker.embed_text(id));
        }
    }
    for &id in &ids[..3] {
        embeds.extend(talker.embed_text(id));
    }
    for (j, crow) in codec_rows[..l - 1].iter().enumerate() {
        let trow = if j < l - 2 { &tts_pad } else { &tts_bos };
        embeds.extend(trow.iter().zip(crow).map(|(a, b)| a + b));
    }

    let trailing_text: Vec<Vec<f32>>;
    match icl {
        None => {
            // First text token summed over codec_bos; the rest streams in per frame.
            let first = talker.embed_text(text_ids[0]);
            embeds.extend(first.iter().zip(&codec_rows[l - 1]).map(|(a, b)| a + b));
            let mut rows: Vec<Vec<f32>> = text_ids[1..]
                .iter()
                .map(|&id| talker.embed_text(id))
                .collect();
            rows.push(tts_eos);
            trailing_text = rows;
        }
        Some((ref_text, ref_codes)) => {
            let rids = tk.encode(&wrap_ref_template(ref_text))?;
            anyhow::ensure!(rids.len() >= 6, "ref text tokenizes too short");
            let ref_ids = &rids[3..rids.len() - 2];
            let mut text_rows: Vec<Vec<f32>> = ref_ids
                .iter()
                .chain(text_ids)
                .map(|&id| talker.embed_text(id))
                .collect();
            text_rows.push(tts_eos);
            let mut icl_codec: Vec<Vec<f32>> = vec![talker.embed_codec(cfg::CODEC_BOS).to_vec()];
            for f in ref_codes {
                let mut sum = talker.embed_codec(f[0]).to_vec();
                for (g, &code) in f[1..].iter().enumerate() {
                    for (s, v) in sum.iter_mut().zip(cp.embed_group(g + 1, code)) {
                        *s += *v;
                    }
                }
                icl_codec.push(sum);
            }
            let (t1, t2) = (text_rows.len(), icl_codec.len());
            if t1 > t2 {
                for (trow, crow) in text_rows[..t2].iter().zip(&icl_codec) {
                    embeds.extend(trow.iter().zip(crow).map(|(a, b)| a + b));
                }
                trailing_text = text_rows[t2..].to_vec();
            } else {
                for (j, crow) in icl_codec.iter().enumerate() {
                    let trow = if j < t1 { &text_rows[j] } else { &tts_pad };
                    embeds.extend(trow.iter().zip(crow).map(|(a, b)| a + b));
                }
                trailing_text = Vec::new();
            }
        }
    }
    Ok(TtsPrompt {
        embeds,
        trailing_text,
        tts_pad,
    })
}