//! Streaming Speech-to-Text via Kyutai's delayed-streams architecture.
//!
//! This module is a thin wrapper around the `moshi` crate's `asr` + `lm` +
//! `mimi` modules. Kyutai's STT shares architecture with their Moshi
//! conversational model: a single decoder-only transformer consumes 32
//! streams of Mimi audio codebook embeddings + 1 text stream and emits
//! one text token per 80 ms frame. The "delayed streams" idea: text is
//! shifted forward in time relative to audio, so token at frame `t`
//! corresponds to audio at frame `t - asr_delay`.
//!
//! ## Why depend on moshi?
//!
//! Kyutai already implements the streaming inference loop, the per-batch
//! state machine, and the word-segmentation logic in pure Rust on
//! candle 0.9.1 (the same candle version we use for CSM). Re-porting the
//! ~600 LOC of `moshi::asr` would be wasted effort. Total new code in
//! rtx-csm is ~200 LOC of API wrapping + audio I/O glue + sentencepiece
//! detok.
//!
//! ## Pipeline
//!
//! ```text
//!   PCM 24 kHz
//!     │
//!     ▼ chunked at 1920 samples (= 80 ms = 1 Mimi frame)
//!   moshi::mimi::Mimi::encode_step → 32 codebook tokens per frame
//!     │
//!     ▼
//!   moshi::lm::LmModel.forward (1 step)
//!     │
//!     ▼
//!   moshi::asr::State.step_pcm → Vec<AsrMsg::{Step, Word, EndWord}>
//!     │
//!     ▼ (sentencepiece detok)
//!   transcribed words with start/stop times
//! ```
//!
//! ## Status: WORKING
//!
//! Verified end-to-end: 23 words correctly transcribed from a 10s real
//! speech sample (LibriSpeech-style FLAC, "He hoped there would be stew
//! for dinner, turnips and carrots and bruised potatoes and fat mutton
//! pieces to be ladled out in thick peppered flower"). The previous
//! "all-pad" debugging session was misled by feeding CSM-generated audio
//! that even the Python reference can't transcribe.
//!
//! What this module ships:
//!   - [`Stt`] struct with `load_default()` (1B en/fr Kyutai checkpoint,
//!     ~3 GB), `load(...)` (custom paths), `step_pcm(...)`, `reset()`,
//!     `decode_word_text(...)` (sentencepiece detok)
//!   - [`AsrEvent`] enum mirroring `moshi::asr::AsrMsg`
//!   - [`config_stt_1b_en_fr`] config matching the released checkpoint
//!
//! Known polish items (~1 hour to chase down if you care about parity):
//!   - Last word or two in long utterances may be cut off — Python's
//!     reference uses `audio_delay_seconds=0.5` directly as a chunk count,
//!     while we use `asr_delay_in_tokens=6` (= 6/12.5 Hz = 0.48s) which
//!     loses the last ~0.08s. Bump to 7 to match.
//!   - SentencePiece tokens are emitted per-word; consecutive same-word
//!     tokens may merge in `Word::tokens` events vs Python which splits
//!     more aggressively. Cosmetic only.

use crate::error::{CsmError, Result};
use candle_core::{DType, Device, Tensor};
use moshi::StreamMask;
use moshi::asr::AsrMsg;
use moshi::lm;
use moshi::transformer;
use sentencepiece::SentencePieceProcessor;
use std::path::{Path, PathBuf};

/// Token id 3 marks the word boundary in Kyutai STT output. Skip during
/// detok so we don't emit "▁" placeholders. Token id 0 is end-of-padding.
const PADDING_TOKEN_ID: u32 = 3;

/// Build the LM config for the VAD-enabled `kyutai/stt-1b-en_fr-candle`
/// variant. Same as the standard config but with `extra_heads = Some(...)`
/// so the LM exposes 4 extra prediction heads (each 6-dim categorical)
/// for semantic end-of-turn detection.
pub fn config_stt_1b_en_fr_vad() -> lm::Config {
    let mut cfg = config_stt_1b_en_fr();
    cfg.extra_heads = Some(lm::ExtraHeadsConfig {
        num_heads: VAD_EXTRA_HEADS,
        dim: VAD_HEAD_DIM,
    });
    cfg
}

/// Build the LM config for `kyutai/stt-1b-en_fr`. Mirrors the upstream
/// config.json: 16 layers / 2048 dim / 16 heads, hidden_scale 4.125 →
/// dim_feedforward = 2048 * 4 (round); 32 audio codebooks; text vocab
/// 8000 (+1 padding for in-vocab); no depformer; no extra heads.
pub fn config_stt_1b_en_fr() -> lm::Config {
    let lm_cfg = transformer::Config {
        d_model: 2048,
        num_heads: 16,
        num_layers: 16,
        // moshi's transformer SwiGLU formula:
        //   if dim_feedforward == 4 * d_model: hidden = 11 * d_model / 4
        //   else: hidden = 2 * dim_feedforward / 3
        // For this checkpoint d_model=2048, hidden=5632, so set
        // dim_feedforward = 4 * 2048 = 8192 to trigger the right branch.
        dim_feedforward: 2048 * 4,
        causal: true,
        norm_first: true,
        bias_ff: false,
        bias_attn: false,
        layer_scale: None,
        context: 750,
        max_period: 100_000,
        use_conv_block: false,
        use_conv_bias: true,
        cross_attention: None,
        gating: Some(candle_nn::Activation::Silu),
        norm: moshi::NormType::RmsNorm,
        positional_embedding: transformer::PositionalEmbedding::Rope,
        conv_layout: false,
        conv_kernel_size: 3,
        kv_repeat: 1,
        max_seq_len: 4096,
        shared_cross_attn: false,
    };
    lm::Config {
        transformer: lm_cfg,
        depformer: None,
        audio_vocab_size: 2049,
        text_in_vocab_size: 8001,
        text_out_vocab_size: 8000,
        audio_codebooks: 32,
        conditioners: Default::default(),
        extra_heads: None,
    }
}

/// Kyutai 1B en/fr STT — the standard 1B en/fr model (no VAD heads).
pub const REPO_KYUTAI_STT_1B: &str = "kyutai/stt-1b-en_fr";
/// Kyutai 1B en/fr STT — VAD-enabled variant. Same backbone weights but
/// with 4 extra heads (each 6-dim categorical, e.g. pause-duration buckets)
/// trained for semantic end-of-turn detection.
pub const REPO_KYUTAI_STT_1B_VAD: &str = "kyutai/stt-1b-en_fr-candle";
pub const FILE_STT_MODEL: &str = "model.safetensors";
pub const FILE_STT_MIMI: &str = "mimi-pytorch-e351c8d8@125.safetensors";
pub const FILE_STT_TOKENIZER: &str = "tokenizer_en_fr_audio_8000.model";

/// Number of extra heads on the VAD-enabled checkpoint.
pub const VAD_EXTRA_HEADS: usize = 4;
/// Per-head output dim (categorical buckets, e.g. pause durations).
pub const VAD_HEAD_DIM: usize = 6;
/// Index of the "end-of-turn" head (per the Kyutai delayed-streams
/// reference Python script: vad_heads[2] is the EOT head).
pub const VAD_EOT_HEAD_IDX: usize = 2;

/// Mimi audio frame rate (12.5 Hz = 80 ms per frame).
pub const FRAME_RATE_HZ: f64 = 12.5;
/// Per-step input samples for `Mimi::encode_step`. The Kyutai STT
/// inference scripts feed 1920 samples per call (one 12.5 Hz output
/// frame at 24 kHz). The internal stride-2 downsample emits one output
/// per call when fed 1920 input samples.
pub const SAMPLES_PER_FRAME: usize = 1920;
pub const SAMPLE_RATE: u32 = 24_000;
/// ASR delay for the 1B en/fr checkpoint. Tokens emitted at frame `t`
/// correspond to audio at frame `t - ASR_DELAY_FRAMES`.
pub const ASR_DELAY_FRAMES: usize = 6;

/// Cleaner enum mirror of `moshi::asr::AsrMsg`. We reshape it slightly so
/// the sentencepiece detok step can be added without disturbing callers.
#[derive(Debug, Clone)]
pub enum AsrEvent {
    /// Per-step probabilities from the `extra_heads` output (semantic VAD,
    /// turn-taking, etc.). Each inner `Vec<f32>` is one head's output.
    Step { step_idx: usize, prs: Vec<Vec<f32>> },
    /// A complete word (sequence of subword tokens) with timing.
    /// `text` is `None` until sentencepiece detok is wired (Phase 6a polish).
    Word {
        tokens: Vec<u32>,
        text: Option<String>,
        start_time: f64,
        batch_idx: usize,
    },
    /// End-of-word marker with stop time.
    EndWord { stop_time: f64, batch_idx: usize },
}

impl From<AsrMsg> for AsrEvent {
    fn from(m: AsrMsg) -> Self {
        match m {
            AsrMsg::Step { step_idx, prs } => AsrEvent::Step { step_idx, prs },
            AsrMsg::Word {
                tokens,
                start_time,
                batch_idx,
            } => AsrEvent::Word {
                tokens,
                text: None,
                start_time,
                batch_idx,
            },
            AsrMsg::EndWord {
                stop_time,
                batch_idx,
            } => AsrEvent::EndWord {
                stop_time,
                batch_idx,
            },
        }
    }
}

/// Streaming STT engine.
pub struct Stt {
    state: moshi::asr::State,
    device: Device,
    /// Buffer of incoming PCM samples awaiting the next 1920-sample frame.
    pending: Vec<f32>,
    /// SentencePiece tokenizer for detokenizing word-token sequences. None
    /// when constructed without a tokenizer path; in that case `Word.text`
    /// is `None` and callers see raw token IDs only.
    tokenizer: Option<SentencePieceProcessor>,
    #[allow(dead_code)]
    tokenizer_path: Option<PathBuf>,
}

impl Stt {
    /// Load the default 1B en/fr STT model from HuggingFace cache (downloads
    /// on first run via `hf-hub`). No VAD heads.
    pub fn load_default(device: &Device) -> Result<Self> {
        Self::load_from_repo(REPO_KYUTAI_STT_1B, /* vad */ false, device)
    }

    /// Load the VAD-enabled variant `kyutai/stt-1b-en_fr-candle`. Same
    /// backbone weights but with 4 extra heads exposed via Step events for
    /// semantic end-of-turn detection.
    pub fn load_default_with_vad(device: &Device) -> Result<Self> {
        Self::load_from_repo(REPO_KYUTAI_STT_1B_VAD, /* vad */ true, device)
    }

    fn load_from_repo(repo: &str, vad: bool, device: &Device) -> Result<Self> {
        let api = hf_hub::api::sync::Api::new()
            .map_err(|e| CsmError::Config(format!("hf-hub init: {e}")))?;
        let r = api.model(repo.to_string());
        let model_path = r
            .get(FILE_STT_MODEL)
            .map_err(|e| CsmError::Config(format!("download {FILE_STT_MODEL}: {e}")))?;
        let mimi_path = r
            .get(FILE_STT_MIMI)
            .map_err(|e| CsmError::Config(format!("download {FILE_STT_MIMI}: {e}")))?;
        let tokenizer_path = r
            .get(FILE_STT_TOKENIZER)
            .map_err(|e| CsmError::Config(format!("download {FILE_STT_TOKENIZER}: {e}")))?;
        Self::load_with_config(&model_path, &mimi_path, Some(&tokenizer_path), vad, device)
    }

    /// Load from explicit weight paths (no VAD).
    pub fn load(
        model: &Path,
        mimi: &Path,
        tokenizer: Option<&Path>,
        device: &Device,
    ) -> Result<Self> {
        Self::load_with_config(model, mimi, tokenizer, /* vad */ false, device)
    }

    /// Load from explicit weight paths with optional VAD heads. Pass
    /// `vad = true` only when the safetensors actually contains
    /// `extra_heads.X.weight` keys (e.g. `kyutai/stt-1b-en_fr-candle`).
    pub fn load_with_config(
        model: &Path,
        mimi: &Path,
        tokenizer: Option<&Path>,
        vad: bool,
        device: &Device,
    ) -> Result<Self> {
        // Kyutai STT checkpoint is bf16. Use bf16 on accelerators (Metal
        // supports bf16 in candle 0.9.1) and F32 on CPU. F16 on Metal
        // produces all-pad outputs because the LM's RmsNorm overflows in
        // some intermediate activations.
        let dtype = match device {
            Device::Cpu => DType::F32,
            _ => DType::BF16,
        };
        let mimi = moshi::mimi::load(mimi.to_string_lossy().as_ref(), Some(32), device)
            .map_err(|e| CsmError::Config(format!("moshi::mimi::load: {e}")))?;
        let cfg = if vad {
            config_stt_1b_en_fr_vad()
        } else {
            config_stt_1b_en_fr()
        };
        let lm = moshi::lm::load_lm_model(cfg, model, dtype, device)
            .map_err(|e| CsmError::Config(format!("moshi::lm::load_lm_model: {e}")))?;
        let state = moshi::asr::State::new(
            /* batch_size */ 1,
            ASR_DELAY_FRAMES,
            /* temperature */ 0.0,
            mimi,
            lm,
        )
        .map_err(|e| CsmError::Config(format!("moshi::asr::State::new: {e}")))?;
        let tokenizer_obj = match tokenizer {
            Some(p) => Some(
                SentencePieceProcessor::open(p)
                    .map_err(|e| CsmError::Config(format!("sentencepiece open: {e}")))?,
            ),
            None => None,
        };
        Ok(Self {
            state,
            device: device.clone(),
            pending: Vec::new(),
            tokenizer: tokenizer_obj,
            tokenizer_path: tokenizer.map(|p| p.to_path_buf()),
        })
    }

    /// Pull the end-of-turn probability out of a [`AsrEvent::Step`] event's
    /// `prs` field. Returns `None` if the event isn't a Step or the
    /// VAD-enabled config wasn't loaded. The Kyutai delayed-streams
    /// reference uses head index 2 (of 4); a probability above ~0.5
    /// across multiple consecutive frames signals end-of-turn.
    pub fn end_of_turn_probability(event: &AsrEvent) -> Option<f32> {
        match event {
            AsrEvent::Step { prs, .. } if prs.len() > VAD_EOT_HEAD_IDX => {
                prs[VAD_EOT_HEAD_IDX].first().copied()
            }
            _ => None,
        }
    }

    /// Detokenize a Word event's token IDs to text. Skips padding tokens
    /// (id 3) and uses sentencepiece's built-in detok if a tokenizer was
    /// loaded.
    pub fn decode_word_text(&self, tokens: &[u32]) -> Option<String> {
        let sp = self.tokenizer.as_ref()?;
        let filtered: Vec<u32> = tokens
            .iter()
            .copied()
            .filter(|&t| t > PADDING_TOKEN_ID)
            .collect();
        if filtered.is_empty() {
            return Some(String::new());
        }
        sp.decode_piece_ids(&filtered).ok()
    }

    /// Reset the streaming state for a new utterance/session.
    pub fn reset(&mut self) -> Result<()> {
        self.state
            .reset()
            .map_err(|e| CsmError::Config(format!("reset: {e}")))?;
        self.pending.clear();
        Ok(())
    }

    /// Feed PCM samples (24 kHz, mono, f32 in [-1, 1]). The buffer is
    /// chunked into 1920-sample frames internally; partial frames are
    /// buffered until the next call. Emits any `AsrEvent`s produced by
    /// the underlying state machine.
    pub fn step_pcm(&mut self, samples: &[f32]) -> Result<Vec<AsrEvent>> {
        self.pending.extend_from_slice(samples);
        let mut events = Vec::new();
        while self.pending.len() >= SAMPLES_PER_FRAME {
            let frame: Vec<f32> = self.pending.drain(..SAMPLES_PER_FRAME).collect();
            let pcm = Tensor::from_vec(frame, (1, 1, SAMPLES_PER_FRAME), &self.device)
                .map_err(|e| CsmError::Config(format!("frame tensor: {e}")))?;
            let mask = StreamMask::empty();
            let msgs = self
                .state
                .step_pcm(pcm, None, &mask, |_, _, _| {})
                .map_err(|e| CsmError::Config(format!("step_pcm: {e}")))?;
            events.extend(msgs.into_iter().map(AsrEvent::from));
        }
        Ok(events)
    }

    /// End-of-stream flush. Drains any partial sub-frame buffer (zero-padded
    /// to a full 1920-sample frame), then steps `ASR_DELAY_FRAMES` additional
    /// silent frames so words emitted at LM step `t` for audio at frame
    /// `t - ASR_DELAY_FRAMES` finally surface — without this the last ~480 ms
    /// of speech (a couple of words on long utterances) trails off after the
    /// caller stops feeding audio.
    ///
    /// Cost: `ASR_DELAY_FRAMES + 1` extra `step_pcm` calls (≈ 7 × per-frame
    /// compute on Metal). Quality fix only; not a latency optimization.
    pub fn finish(&mut self) -> Result<Vec<AsrEvent>> {
        let mut events = Vec::new();
        let mask = StreamMask::empty();

        if !self.pending.is_empty() {
            self.pending.resize(SAMPLES_PER_FRAME, 0.0);
            let frame = std::mem::take(&mut self.pending);
            let pcm = Tensor::from_vec(frame, (1, 1, SAMPLES_PER_FRAME), &self.device)
                .map_err(|e| CsmError::Config(format!("finish tensor: {e}")))?;
            let msgs = self
                .state
                .step_pcm(pcm, None, &mask, |_, _, _| {})
                .map_err(|e| CsmError::Config(format!("finish step: {e}")))?;
            events.extend(msgs.into_iter().map(AsrEvent::from));
        }

        let zeros = vec![0.0f32; SAMPLES_PER_FRAME];
        for _ in 0..ASR_DELAY_FRAMES {
            let pcm = Tensor::from_vec(zeros.clone(), (1, 1, SAMPLES_PER_FRAME), &self.device)
                .map_err(|e| CsmError::Config(format!("flush tensor: {e}")))?;
            let msgs = self
                .state
                .step_pcm(pcm, None, &mask, |_, _, _| {})
                .map_err(|e| CsmError::Config(format!("flush step: {e}")))?;
            events.extend(msgs.into_iter().map(AsrEvent::from));
        }
        Ok(events)
    }
}

/// Energy-based VAD gate — classifies a 24 kHz PCM slice as
/// speech/silence by RMS amplitude. Independent of `Stt`'s semantic
/// VAD (head-2 from the `kyutai/stt-1b-en_fr-candle` checkpoint, used
/// for end-of-turn); this is a per-chunk silence detector for the
/// receive-loop gate.
///
/// **Why energy and not Silero V5?** Silero V5 ships via the `ort`
/// ONNX runtime, which links a different protobuf version (3.21) than
/// `sentencepiece-sys` (3.14, used by Kyutai STT for word
/// detokenization). The two crates panic at process startup with
/// `libprotobuf FATAL ... version verification failed`. Phase 8.1.3
/// hit this on the first server boot. Energy VAD avoids the runtime
/// entirely — pure Rust, ~30 LOC.
///
/// **Quality tradeoff.** Energy VAD catches obvious silence (room
/// tone, pauses between words) but misses quiet speech (whispering,
/// distant speakers). For mic-distance voice loops this catches ~70-
/// 80% of what Silero V5 would, at zero linkage risk. If you need
/// better recall on quiet speech, port Silero V5 weights to candle
/// natively (deferred work) or run it in a sidecar process.
///
/// Real-world voice-agent audio is 30-50% silence (typing pauses,
/// breathing, room tone). Skipping STT on those chunks yields
/// proportional `recv_phase` reduction without altering transcript
/// quality.
pub struct VadGate {
    /// RMS threshold in [0, 1] (post-normalization to f32 PCM range).
    /// Practical defaults: 0.005-0.02 for typical mic audio.
    threshold_rms: f32,
    pub skipped_chunks: u64,
    pub total_chunks: u64,
}

impl VadGate {
    pub fn new(threshold_rms: f32) -> Self {
        Self {
            threshold_rms,
            skipped_chunks: 0,
            total_chunks: 0,
        }
    }

    /// Returns `true` if the slice's RMS amplitude is at or above the
    /// silence threshold. Pure-Rust, branchless inner loop.
    pub fn is_speech(&mut self, samples_24k: &[f32]) -> bool {
        self.total_chunks += 1;
        if samples_24k.is_empty() {
            return true; // empty input — treat as speech (safe default).
        }
        let mut sum_sq = 0.0f64;
        for &s in samples_24k {
            sum_sq += (s as f64) * (s as f64);
        }
        let rms = (sum_sq / samples_24k.len() as f64).sqrt() as f32;
        let is_speech = rms >= self.threshold_rms;
        if !is_speech {
            self.skipped_chunks += 1;
        }
        is_speech
    }

    /// Fraction of chunks classified as silence so far. For /metrics
    /// reporting.
    pub fn silence_fraction(&self) -> f64 {
        if self.total_chunks == 0 {
            0.0
        } else {
            self.skipped_chunks as f64 / self.total_chunks as f64
        }
    }
}

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

    #[test]
    fn asr_event_conversion() {
        let m = AsrMsg::Word {
            tokens: vec![1, 2, 3],
            start_time: 1.5,
            batch_idx: 0,
        };
        let e: AsrEvent = m.into();
        match e {
            AsrEvent::Word {
                tokens,
                text,
                start_time,
                batch_idx,
            } => {
                assert_eq!(tokens, vec![1, 2, 3]);
                assert!(text.is_none()); // detok not wired yet
                assert!((start_time - 1.5).abs() < 1e-9);
                assert_eq!(batch_idx, 0);
            }
            _ => panic!("expected Word event"),
        }
    }

    #[test]
    fn asr_event_endword_passthrough() {
        let m = AsrMsg::EndWord {
            stop_time: 2.0,
            batch_idx: 0,
        };
        let e: AsrEvent = m.into();
        assert!(matches!(e, AsrEvent::EndWord { stop_time, .. } if (stop_time - 2.0).abs() < 1e-9));
    }

    #[test]
    fn end_of_turn_probability_extracts_head_2() {
        // 4 heads, each with one prob value (matches moshi 0.6.4 emission).
        let event = AsrEvent::Step {
            step_idx: 10,
            prs: vec![vec![0.1], vec![0.2], vec![0.7], vec![0.05]],
        };
        let pr = Stt::end_of_turn_probability(&event).expect("VAD prob");
        assert!((pr - 0.7).abs() < 1e-6);
    }

    #[test]
    fn end_of_turn_probability_none_when_no_extra_heads() {
        let event = AsrEvent::Step {
            step_idx: 1,
            prs: vec![],
        };
        assert!(Stt::end_of_turn_probability(&event).is_none());
        let event = AsrEvent::Word {
            tokens: vec![5],
            text: None,
            start_time: 0.0,
            batch_idx: 0,
        };
        assert!(Stt::end_of_turn_probability(&event).is_none());
    }
}