//! High-level Generator façade.

use crate::audio_io::TARGET_SAMPLE_RATE;
use crate::config::ModelConfig;
use crate::error::Result;
use crate::hub;
use crate::mimi::Mimi;
use crate::model::CsmModel;
use crate::post::PostProcess;
use crate::prompt::{Segment, build_prompt};
use crate::repetition::{RepetitionConfig, RepetitionGuard};
use crate::sampler::{CsmSampler, DEFAULT_TEMPERATURE, DEFAULT_TOPK, DEFAULT_TOPP};
use crate::text_norm::TextNormalize;
use crate::tokenizer::CsmTokenizer;
use crate::util;
use candle_core::{DType, Device, Tensor};

#[derive(Debug, Clone)]
pub struct GenerateOptions {
    pub max_audio_ms: u32,
    pub temperature: f64,
    pub top_k: usize,
    /// Nucleus (top-p) filter applied after top-k. Set to `1.0` (or `0.0`) to disable.
    pub top_p: f64,
    pub seed: u64,
    /// Loop-escape guard. `None` disables (not recommended).
    pub repetition: Option<RepetitionConfig>,
    /// Classifier-Free Guidance scale on the codebook-0 head. `None` or `Some(1.0)`
    /// disables CFG. Koel-TTS recommends 1.5–3.0. Requires the generator to
    /// have been loaded via [`Generator::load_csm_1b_with_cfg`] AND the
    /// generation call to provide non-empty context (without context the
    /// uncond branch == cond branch and CFG has no effect).
    ///
    /// Equivalent to `cfg_schedule = Some(CfgSchedule::Constant(s))` when
    /// `cfg_schedule` is None. The schedule field takes precedence when set.
    pub cfg_scale: Option<f64>,
    /// Optional per-frame CFG scale schedule (Selective CFG, arXiv 2509.19668).
    /// Overrides `cfg_scale` when set. Use `Step { early, late, transition }`
    /// to start at one CFG scale for the first `transition` frames and drop
    /// to a lower one — preserves text adherence late in synthesis while
    /// keeping speaker fidelity early.
    pub cfg_schedule: Option<crate::cfg_schedule::CfgSchedule>,
    /// Optional control-token hint prepended to the text prompt, e.g.
    /// `Some("[whisper]".into())`. Concatenated as `<hint> <text>` and fed
    /// through the same Llama BPE path as ordinary text — the hint is plain
    /// text from the model's perspective. Only useful when the LoRA / fine-tune
    /// has been trained with the same tags in its prompts (per
    /// `personal_voice_training_guide.md` §4): the adapter learns
    /// `[whisper] <text>` → whispered prosody. Out of the box (no fine-tune)
    /// this is a no-op cosmetic prefix.
    pub emotion_hint: Option<String>,
}

impl Default for GenerateOptions {
    fn default() -> Self {
        Self {
            max_audio_ms: 10_000,
            temperature: DEFAULT_TEMPERATURE,
            top_k: DEFAULT_TOPK,
            top_p: DEFAULT_TOPP,
            seed: 42,
            repetition: Some(RepetitionConfig::default()),
            cfg_scale: None,
            cfg_schedule: None,
            emotion_hint: None,
        }
    }
}

/// Prepend an optional control-token hint to the (already-normalized) text.
/// Joined with a single space when present so the BPE tokenizer treats the tag
/// as a separate sub-sequence from the body text. Centralized here so every
/// generate variant — and the LoRA trainer in `training.rs` — applies the
/// same convention. Training and inference *must* use identical prefix
/// formatting or the adapter won't transfer.
pub(crate) fn apply_emotion_hint(text: String, hint: Option<&str>) -> String {
    match hint {
        Some(tag) if !tag.is_empty() => format!("{} {}", tag.trim(), text),
        _ => text,
    }
}

pub struct Generator {
    pub model: CsmModel,
    pub mimi: Mimi,
    pub tokenizer: CsmTokenizer,
    pub config: ModelConfig,
    pub device: Device,
    /// Applied to every text input before tokenization. Set to
    /// [`TextNormalize::passthrough`] if you've pre-normalized upstream.
    pub text_normalize: TextNormalize,
    /// Optional watermarker applied inside [`Self::generate_to_wav`] AFTER
    /// post-processing and BEFORE WAV write. Set via
    /// [`Self::set_watermarker`]. `None` = no-op (Sesame's reference TTS
    /// also ships unwatermarked by default; this is the integration hook).
    pub watermarker: Option<Box<dyn crate::watermark::Watermarker>>,
}

impl Generator {
    pub fn new(model: CsmModel, mimi: Mimi, tokenizer: CsmTokenizer, device: Device) -> Self {
        let config = model.config.clone();
        Self {
            model,
            mimi,
            tokenizer,
            config,
            device,
            text_normalize: TextNormalize::default(),
            watermarker: None,
        }
    }

    /// Install a watermarker that runs inside `generate_to_wav` after
    /// post-processing. Take ownership of the watermarker so the generator
    /// can be moved into worker threads (Watermarker is `Send + Sync`).
    pub fn set_watermarker(&mut self, wm: Box<dyn crate::watermark::Watermarker>) {
        self.watermarker = Some(wm);
    }

    /// Drop any installed watermarker.
    pub fn clear_watermarker(&mut self) {
        self.watermarker = None;
    }

    /// Like [`Self::load_csm_1b`] but loads the CSM weights from an
    /// explicit path instead of the HuggingFace cache. Mimi + Llama
    /// tokenizer are still resolved through the hub. Use this with a
    /// LoRA-merged checkpoint produced by
    /// `lora::merge_into_safetensors`.
    pub fn load_csm_1b_from_path<P: AsRef<std::path::Path>>(
        csm_weights: P,
        device: &Device,
    ) -> Result<Self> {
        let assets = hub::resolve_csm_1b()?;
        let csm_weights = csm_weights.as_ref().to_path_buf();
        let config = ModelConfig::csm_1b();
        let dtype = match device {
            Device::Cpu => DType::F32,
            Device::Metal(_) => DType::F16,
            _ => DType::BF16,
        };
        let model = CsmModel::load_from_safetensors(&csm_weights, config.clone(), dtype, device)?;
        let mimi = Mimi::load(&assets.mimi_weights, device)?;
        let tokenizer = CsmTokenizer::from_file(&assets.tokenizer_json)?;
        Ok(Self::new(model, mimi, tokenizer, device.clone()))
    }

    /// Download (cached) CSM-1B + Mimi + Llama tokenizer from HuggingFace and
    /// build a ready-to-generate `Generator`.
    pub fn load_csm_1b(device: &Device) -> Result<Self> {
        let assets = hub::resolve_csm_1b()?;
        let config = ModelConfig::csm_1b();

        // Optional: audit tensor keys once before loading to surface naming drift.
        if std::env::var("CSM_AUDIT_KEYS").is_ok() {
            let descs = crate::model::dump_safetensors_keys(&assets.csm_weights)?;
            let missing = crate::model::audit_csm_keys(&descs);
            if missing.is_empty() {
                tracing::info!("safetensors key audit: ok ({} tensors)", descs.len());
            } else {
                tracing::warn!("safetensors missing keys: {missing:?}");
            }
        }

        // dtype selection by backend:
        //   CPU    → F32  (candle's CPU backend has no BF16 matmul kernel)
        //   Metal  → F16  (BF16 is ~50% slower than F16 on M1/M2; M3+ added hw bf16
        //                  but f16 still ties or wins. CSM softmax is well-behaved
        //                  post-RMSNorm so f16 dynamic range is fine in practice.)
        //   CUDA   → BF16 (modern NVIDIA tensor cores prefer BF16)
        let dtype = match device {
            Device::Cpu => DType::F32,
            Device::Metal(_) => DType::F16,
            _ => DType::BF16,
        };
        let model =
            CsmModel::load_from_safetensors(&assets.csm_weights, config.clone(), dtype, device)?;
        let mimi = Mimi::load(&assets.mimi_weights, device)?;
        let tokenizer = CsmTokenizer::from_file(&assets.tokenizer_json)?;

        Ok(Self::new(model, mimi, tokenizer, device.clone()))
    }

    /// Like [`Self::load_csm_1b`] but allocates a second backbone for the
    /// unconditional CFG branch. Adds ~70 MB of KV cache; weight tensors are
    /// shared via mmap.
    pub fn load_csm_1b_with_cfg(device: &Device, enable_cfg: bool) -> Result<Self> {
        let assets = hub::resolve_csm_1b()?;
        let config = ModelConfig::csm_1b();
        let dtype = match device {
            Device::Cpu => DType::F32,
            Device::Metal(_) => DType::F16,
            _ => DType::BF16,
        };
        let model = CsmModel::load_from_safetensors_with_cfg(
            &assets.csm_weights,
            config.clone(),
            dtype,
            device,
            enable_cfg,
        )?;
        let mimi = Mimi::load(&assets.mimi_weights, device)?;
        let tokenizer = CsmTokenizer::from_file(&assets.tokenizer_json)?;
        Ok(Self::new(model, mimi, tokenizer, device.clone()))
    }

    /// Load a quantized CSM-1B from a GGUF file (the artifact emitted by
    /// `examples/quantize`). Defaults to candle's raw QMatMul kernel path —
    /// Q8 weights stay quantized at runtime, the Metal/CPU kernel performs
    /// fused dequant+matmul. This is the actual quantization perf win:
    /// smaller weight memory at runtime AND correct, well-pronounced output.
    ///
    /// Verified 2026-04-25 via per-layer hidden-state bisection: hidden
    /// states match between raw-QTensor and F16-dequant paths to <0.02% on
    /// every backbone layer. Earlier "gibberish" was traced to the v1 GGUF
    /// having an incomplete quant policy that left MLP weights at F16 while
    /// attention was Q8 — the mixed-dtype model interacted with the kernel
    /// paths. The kernel itself is correct.
    ///
    /// To force F16-dequantization-on-load (occasionally useful for
    /// numerical-precision debugging or to test against non-Q kernels), set
    /// `CSM_DEQUANT_F16=1` in the environment before invoking. That path
    /// gives F16 runtime weights — same memory profile as the F16 safetensors
    /// path, no inference speedup.
    pub fn load_csm_1b_quantized<P: AsRef<std::path::Path>>(
        gguf_path: P,
        device: &Device,
        enable_cfg: bool,
    ) -> Result<Self> {
        // Opt-in fallback: F16 dequantization at load time.
        if std::env::var("CSM_DEQUANT_F16").is_ok() {
            // SAFETY: candle reads this env var via a thread-local on first
            // access. Setting before QMatMul construction is the supported
            // pattern.
            unsafe { std::env::set_var("CANDLE_DEQUANTIZE_ALL_F16", "1") };
        }

        let mimi_weights = hub::resolve_mimi()?;
        let tokenizer_json = hub::resolve_llama_tokenizer()?;
        let config = ModelConfig::csm_1b();
        // F32 runtime is the safest choice for the quantized path even with
        // DEQUANTIZE_ALL_F16. Reason: tensors stored as F16 in the GGUF
        // (kept-native heads/embeds/MLPs etc.) get auto-dequantized to F16
        // tensors by candle's QMatMul construction, which would clash with
        // F16 runtime when also mixing TensorF16-wrapped Q8 weights. Routing
        // everything through F32 activations + F16 dequantize_f16 paths
        // (auto-cast inside TensorF16 forward) avoids any dtype mismatch.
        let runtime_dtype = DType::F32;
        let model =
            CsmModel::load_from_gguf(gguf_path, config.clone(), runtime_dtype, device, enable_cfg)?;
        let mimi = Mimi::load(&mimi_weights, device)?;
        let tokenizer = CsmTokenizer::from_file(&tokenizer_json)?;
        Ok(Self::new(model, mimi, tokenizer, device.clone()))
    }

    pub fn reset(&mut self) {
        self.model.clear_kv_cache();
    }

    /// Install activation-steering vectors on the conditional Llama backbone.
    /// Only the FP backend supports steering today; on a quantized backend
    /// this returns an error rather than silently no-op'ing. See
    /// `crate::steering`. Pass `None` to remove.
    pub fn set_steering(
        &mut self,
        steering: Option<crate::steering::LayerSteering>,
    ) -> crate::Result<()> {
        match &mut self.model.inner {
            crate::model::ModelBackend::Fp(m) => {
                m.set_backbone_steering(steering);
                Ok(())
            }
            crate::model::ModelBackend::Quantized(_) => Err(crate::error::CsmError::Config(
                "activation steering is only supported on the FP backbone \
                 (--quantized-gguf disables it)"
                    .into(),
            )),
        }
    }

    /// Number of backbone layers — useful when constructing
    /// `LayerSteering::empty(n)` to match the loaded model. Returns `None`
    /// for the quantized backend (steering not supported there).
    pub fn backbone_num_layers(&self) -> Option<usize> {
        match &self.model.inner {
            crate::model::ModelBackend::Fp(m) => Some(m.backbone_num_layers()),
            crate::model::ModelBackend::Quantized(_) => None,
        }
    }

    /// Install steering on the depth decoder (acoustic codebooks). Vectors
    /// must be `(decoder_embed_dim,)` — for CSM-1B that's 1024, NOT 2048.
    /// Architectural hypothesis: backbone carries semantic content (what
    /// the model says), decoder carries acoustic detail (how it sounds);
    /// steering the decoder should shift prosody/timbre without disturbing
    /// word-level fidelity.
    pub fn set_decoder_steering(
        &mut self,
        steering: Option<crate::steering::LayerSteering>,
    ) -> crate::Result<()> {
        match &mut self.model.inner {
            crate::model::ModelBackend::Fp(m) => {
                m.set_decoder_steering(steering);
                Ok(())
            }
            crate::model::ModelBackend::Quantized(_) => Err(crate::error::CsmError::Config(
                "decoder steering only supported on the FP backbone".into(),
            )),
        }
    }

    /// Number of decoder layers (CSM-1B = 4). `None` on quantized.
    pub fn decoder_num_layers(&self) -> Option<usize> {
        match &self.model.inner {
            crate::model::ModelBackend::Fp(m) => Some(m.decoder_num_layers()),
            crate::model::ModelBackend::Quantized(_) => None,
        }
    }

    /// Full generation loop: prompt → backbone/decoder per-frame → Mimi decode → f32 PCM.
    pub fn generate(
        &mut self,
        text: &str,
        speaker: u32,
        context: &[Segment],
        opts: GenerateOptions,
    ) -> Result<Vec<f32>> {
        self.reset();

        let normalized = self.text_normalize.apply(text)?;
        let prompted = apply_emotion_hint(normalized, opts.emotion_hint.as_deref());
        let current = Segment::new_text(speaker, prompted);
        let prompt = build_prompt(
            context,
            &current,
            &self.model,
            &mut self.mimi,
            &self.tokenizer,
        )?;

        let cb = self.config.audio_num_codebooks;
        let mut sampler = CsmSampler::new(opts.seed, opts.temperature, opts.top_k, opts.top_p);
        let inner_lp = sampler.inner_mut();

        // CFG path: dual backbone if all preconditions are met. The
        // schedule field takes precedence over the legacy `cfg_scale`.
        let cfg_schedule = opts.cfg_schedule.or_else(|| {
            opts.cfg_scale
                .map(crate::cfg_schedule::CfgSchedule::Constant)
        });
        let cfg_active = cfg_schedule.map(|s| s.is_active()).unwrap_or(false)
            && !context.is_empty()
            && self.model.inner.cfg_enabled();
        let cfg_schedule = if cfg_active {
            cfg_schedule.unwrap()
        } else {
            crate::cfg_schedule::CfgSchedule::Constant(1.0)
        };

        let uncond_prompt_opt = if cfg_active {
            Some(build_prompt(
                &[],
                &current,
                &self.model,
                &mut self.mimi,
                &self.tokenizer,
            )?)
        } else {
            None
        };

        let mut pos: usize = 0;
        let mut uncond_pos: usize = 0;
        let mut all_frames: Vec<Vec<u32>> = Vec::new();
        let max_frames =
            ((opts.max_audio_ms as f32) / self.config.frame_duration_ms()).ceil() as usize;

        let mut input_tokens = prompt.tokens;
        let mut input_mask = prompt.mask;
        let mut uncond_tokens = uncond_prompt_opt.as_ref().map(|p| p.tokens.clone());
        let mut uncond_mask = uncond_prompt_opt.as_ref().map(|p| p.mask.clone());

        let mut rep_guard = opts.repetition.map(RepetitionGuard::new);

        if cfg_active {
            tracing::info!("CFG active (schedule={cfg_schedule:?})");
        }

        for frame_idx in 0..max_frames {
            let sampled = if cfg_active {
                let ut = uncond_tokens.as_ref().unwrap();
                let um = uncond_mask.as_ref().unwrap();
                let frame_scale = cfg_schedule.scale_at(frame_idx);
                let r = self.model.inner.generate_frame_cfg(
                    &input_tokens,
                    &input_mask,
                    pos,
                    ut,
                    um,
                    uncond_pos,
                    frame_scale,
                    inner_lp,
                )?;
                uncond_pos += ut.dim(1)?;
                r
            } else {
                self.model
                    .inner
                    .generate_frame(&input_tokens, &input_mask, pos, inner_lp)?
            };
            pos += input_tokens.dim(1)?;

            let is_eot = frame_idx >= 1 && sampled.iter().all(|v| *v == 0);
            if is_eot {
                tracing::info!("EOT detected at frame {frame_idx}");
                break;
            }

            if let Some(g) = rep_guard.as_mut()
                && g.observe(&sampled)
            {
                tracing::warn!(
                    "loop-escape: repetition guard tripped at frame {frame_idx}; ending generation"
                );
                break;
            }

            all_frames.push(sampled.clone());

            let (t, m) = self.model.inner.audio_tokens_and_mask(sampled)?;
            // Both branches feed the same sampled frame (sampling produces a single
            // discrete decision; the uncond cache must track the chosen path too).
            if cfg_active {
                uncond_tokens = Some(t.clone());
                uncond_mask = Some(m.clone());
            }
            input_tokens = t;
            input_mask = m;
        }

        if all_frames.is_empty() {
            return Ok(Vec::new());
        }

        // Pack frames into (1, cb, T) i64 and decode through Mimi.
        let t = all_frames.len();
        let mut flat: Vec<u32> = Vec::with_capacity(t * cb);
        // Transpose: all_frames is Vec<frame: Vec<cb>>; we want row-major (cb, t).
        for c in 0..cb {
            for frame in &all_frames {
                flat.push(frame[c]);
            }
        }
        let codes = Tensor::from_vec(flat, (1, cb, t), &self.device)?.to_dtype(DType::U32)?;
        let pcm = self.mimi.decode(&codes)?;
        tracing::info!(
            "generated {} samples (~{:.2}s at {} Hz)",
            pcm.len(),
            pcm.len() as f32 / TARGET_SAMPLE_RATE as f32,
            TARGET_SAMPLE_RATE
        );
        Ok(pcm)
    }

    /// Streaming version of [`Self::generate`]: invokes `on_chunk` with new PCM
    /// samples every `chunk_frames` frames (default 4 = ~320 ms) so a downstream
    /// player can start audio output before generation completes.
    ///
    /// Uses Mimi's `decode_step` (StreamTensor-based incremental decode), so
    /// total decode cost is O(n) instead of O(n²). First-audio latency drops
    /// from ~T_total to `chunk_frames × 80 ms + per_frame_compute × chunk_frames`.
    pub fn generate_streaming<F>(
        &mut self,
        text: &str,
        speaker: u32,
        context: &[Segment],
        opts: GenerateOptions,
        chunk_frames: usize,
        mut on_chunk: F,
    ) -> Result<Vec<f32>>
    where
        F: FnMut(&[f32]) -> Result<()>,
    {
        self.reset();
        // Mimi's streaming state must be reset at the start of every stream;
        // leftover state from a prior call corrupts the first chunk.
        self.mimi.reset_state();

        let chunk_frames = chunk_frames.max(1);

        let normalized = self.text_normalize.apply(text)?;
        let prompted = apply_emotion_hint(normalized, opts.emotion_hint.as_deref());
        let current = Segment::new_text(speaker, prompted);
        let prompt = build_prompt(
            context,
            &current,
            &self.model,
            &mut self.mimi,
            &self.tokenizer,
        )?;

        let cb = self.config.audio_num_codebooks;
        let mut sampler = CsmSampler::new(opts.seed, opts.temperature, opts.top_k, opts.top_p);
        let inner_lp = sampler.inner_mut();

        let mut pos: usize = 0;
        let max_frames =
            ((opts.max_audio_ms as f32) / self.config.frame_duration_ms()).ceil() as usize;
        let mut input_tokens = prompt.tokens;
        let mut input_mask = prompt.mask;

        // Pending frames not yet sent to Mimi's decode_step.
        let mut pending: Vec<Vec<u32>> = Vec::with_capacity(chunk_frames);
        let mut full_pcm: Vec<f32> = Vec::new();
        let mut rep_guard = opts.repetition.map(RepetitionGuard::new);

        for frame_idx in 0..max_frames {
            let sampled =
                self.model
                    .inner
                    .generate_frame(&input_tokens, &input_mask, pos, inner_lp)?;
            pos += input_tokens.dim(1)?;

            let is_eot = frame_idx >= 1 && sampled.iter().all(|v| *v == 0);
            if is_eot {
                tracing::info!("EOT detected at frame {frame_idx}");
                break;
            }
            if let Some(g) = rep_guard.as_mut()
                && g.observe(&sampled)
            {
                tracing::warn!(
                    "loop-escape: repetition guard tripped at frame {frame_idx}; ending"
                );
                break;
            }

            pending.push(sampled.clone());

            if pending.len() >= chunk_frames {
                stream_decode_pending(
                    &mut pending,
                    cb,
                    &mut self.mimi,
                    &self.device,
                    &mut full_pcm,
                    &mut on_chunk,
                )?;
            }

            let (t, m) = self.model.inner.audio_tokens_and_mask(sampled)?;
            input_tokens = t;
            input_mask = m;
        }

        // Flush any frames left in the buffer.
        if !pending.is_empty() {
            stream_decode_pending(
                &mut pending,
                cb,
                &mut self.mimi,
                &self.device,
                &mut full_pcm,
                &mut on_chunk,
            )?;
        }
        // Drain any internal Mimi buffering by feeding a `None` step.
        if let Some(tail) = self.mimi.decode_step(None)?
            && !tail.is_empty()
        {
            on_chunk(&tail)?;
            full_pcm.extend_from_slice(&tail);
        }

        tracing::info!(
            "streaming generation finished: {} samples (~{:.2}s)",
            full_pcm.len(),
            full_pcm.len() as f32 / self.config.sample_rate as f32
        );
        Ok(full_pcm)
    }

    /// Generate using a `SpeakerProfile` as context. The profile is automatically
    /// fit to budget; the caller's provided context segments come *after* the profile.
    pub fn generate_with_profile(
        &mut self,
        profile: &crate::speaker::SpeakerProfile,
        text: &str,
        extra_context: &[Segment],
        opts: GenerateOptions,
    ) -> Result<Vec<f32>> {
        let mut profile = profile.clone();
        profile.fit_within_budget(crate::speaker::DEFAULT_PROFILE_BUDGET_TOKENS);
        let mut ctx = profile.segments().to_vec();
        ctx.extend_from_slice(extra_context);
        self.generate(text, profile.id, &ctx, opts)
    }

    /// Convenience: `generate` + apply default post-processing + watermark
    /// (if installed) + write WAV. Pass [`PostProcess::disabled`] to skip
    /// post-processing. Pass [`Self::clear_watermarker`] (or never install
    /// one) to skip watermarking.
    ///
    /// Order: model → post-process (HPF + declick + LUFS) → watermark.
    /// Watermarking comes last so the loudness target the user sees on disk
    /// is the loudness target the user requested (the watermark residual
    /// is at most a few dB and well below LUFS measurement floor).
    pub fn generate_to_wav(
        &mut self,
        text: &str,
        speaker: u32,
        context: &[Segment],
        opts: GenerateOptions,
        post: &PostProcess,
        out_path: &std::path::Path,
    ) -> Result<()> {
        let mut pcm = self.generate(text, speaker, context, opts)?;
        post.apply(&mut pcm, self.config.sample_rate)?;
        if let Some(wm) = self.watermarker.as_ref() {
            pcm = wm.embed(&pcm)?;
        }
        crate::audio_io::write_wav_24k_mono(out_path, &pcm)?;
        Ok(())
    }

    /// Choose the best available device based on enabled features.
    pub fn default_device() -> Result<Device> {
        util::pick_device()
    }
}

/// Take all `pending` frames, build a (1, cb, k) tensor, push to Mimi's
/// streaming decoder, emit any audio it produces, and clear `pending`.
fn stream_decode_pending<F>(
    pending: &mut Vec<Vec<u32>>,
    num_codebooks: usize,
    mimi: &mut Mimi,
    device: &Device,
    full_pcm: &mut Vec<f32>,
    on_chunk: &mut F,
) -> Result<()>
where
    F: FnMut(&[f32]) -> Result<()>,
{
    if pending.is_empty() {
        return Ok(());
    }
    let k = pending.len();
    let mut flat: Vec<u32> = Vec::with_capacity(k * num_codebooks);
    for c in 0..num_codebooks {
        for frame in pending.iter() {
            flat.push(frame[c]);
        }
    }
    let codes = Tensor::from_vec(flat, (1, num_codebooks, k), device)?.to_dtype(DType::U32)?;
    pending.clear();
    if let Some(samples) = mimi.decode_step(Some(&codes))?
        && !samples.is_empty()
    {
        on_chunk(&samples)?;
        full_pcm.extend_from_slice(&samples);
    }
    Ok(())
}

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

    #[test]
    fn apply_emotion_hint_no_op_when_none() {
        let out = apply_emotion_hint("Hello there".into(), None);
        assert_eq!(out, "Hello there");
    }

    #[test]
    fn apply_emotion_hint_no_op_when_empty() {
        let out = apply_emotion_hint("Hello there".into(), Some(""));
        assert_eq!(out, "Hello there");
    }

    #[test]
    fn apply_emotion_hint_prepends_with_space() {
        let out = apply_emotion_hint("Hello there".into(), Some("[whisper]"));
        assert_eq!(out, "[whisper] Hello there");
    }

    #[test]
    fn apply_emotion_hint_trims_outer_whitespace_on_tag() {
        // Stray whitespace on the tag itself should not double up the separator.
        let out = apply_emotion_hint("Hi".into(), Some("  [excited]  "));
        assert_eq!(out, "[excited] Hi");
    }
}