decibri 3.1.0

//! Voice Activity Detection using the Silero VAD v5 ONNX model.
//!
//! Silero VAD is a pre-trained ML model that detects human speech in audio.
//! It runs inference on 512-sample windows (at 16kHz) and returns a speech
//! probability between 0.0 and 1.0.
//!
//! # Model tensor specification (verified from silero_vad.onnx)
//!
//! Inputs:
//!   - `input`:  f32[batch, window_size]  (audio samples)
//!   - `state`:  f32[2, batch, 128]       (LSTM state, hidden + cell combined)
//!   - `sr`:     i64 scalar               (sample rate)
//!
//! Outputs:
//!   - `output`: f32[batch, 1]            (speech probability)
//!   - `stateN`: f32[2, batch, 128]       (updated LSTM state)

#[cfg(feature = "vad")]
use std::path::{Path, PathBuf};
#[cfg(feature = "vad")]
use std::sync::OnceLock;

use crate::error::DecibriError;

/// Configuration for Silero VAD.
#[derive(Debug, Clone)]
pub struct VadConfig {
    /// Path to the silero_vad.onnx model file.
    pub model_path: PathBuf,
    /// Sample rate: 8000 or 16000 Hz.
    pub sample_rate: u32,
    /// Speech probability threshold (0.0 to 1.0). Default: 0.5.
    pub threshold: f32,
    /// Optional absolute path to the ONNX Runtime shared library.
    ///
    /// When `Some(path)`, ORT is initialized from the given library path via
    /// `ort::init_from`. When `None`, `ort::init()` is called and ORT honours
    /// the `ORT_DYLIB_PATH` environment variable if set.
    ///
    /// ORT is initialized exactly once per process. If multiple `SileroVad`
    /// instances are constructed with different `ort_library_path` values,
    /// the first path wins and later paths are silently ignored. This matches
    /// ORT's own single-global-runtime model.
    pub ort_library_path: Option<PathBuf>,
}

impl Default for VadConfig {
    fn default() -> Self {
        Self {
            model_path: PathBuf::from("silero_vad.onnx"),
            sample_rate: 16000,
            threshold: 0.5,
            ort_library_path: None,
        }
    }
}

/// Result of processing an audio chunk through Silero VAD.
#[derive(Debug, Clone)]
pub struct VadResult {
    /// Maximum speech probability across all windows in the chunk.
    pub probability: f32,
    /// Whether probability >= threshold.
    pub is_speech: bool,
}

/// Silero VAD v5 inference engine.
///
/// Stateful: maintains LSTM hidden/cell state across calls.
/// Call `process()` with each audio chunk. Call `reset()` to clear state.
#[cfg(feature = "vad")]
pub struct SileroVad {
    session: ort::session::Session,
    /// Combined LSTM state [2, 1, 128] = 256 floats.
    state: Vec<f32>,
    sample_rate: u32,
    threshold: f32,
    /// Carries leftover samples between process() calls.
    accumulator: Vec<f32>,
    /// 512 for 16kHz, 256 for 8kHz.
    window_size: usize,
}

/// State size: 2 (hidden + cell) * 1 (batch) * 128 (hidden_dim) = 256.
/// State size: 2 (hidden + cell layers) × batch(1) × hidden_dim(128).
const STATE_SIZE: usize = 256;

/// Process-global ORT init tracker. Stores the library path used on first
/// successful init (None if init used `ort::init()` with env-var fallback).
/// Subsequent `init_ort_once` calls see this as populated and return immediately.
#[cfg(feature = "vad")]
static ORT_INIT: OnceLock<Option<PathBuf>> = OnceLock::new();

/// Wrap an ORT init failure with a decibri-specific actionable message.
///
/// Kept as a standalone generic function (rather than an inline closure in
/// `init_ort_once`) so it can be unit-tested without triggering a real ORT
/// init failure. Tests pass a synthetic `err` value.
#[cfg(feature = "vad")]
fn wrap_init_error<E: std::fmt::Display>(path: Option<&Path>, err: E) -> DecibriError {
    match path {
        Some(p) => DecibriError::Other(format!(
            "decibri: failed to load ONNX Runtime from {}: {}. \
             If ORT_DYLIB_PATH is set, verify it points to a valid ONNX Runtime \
             library for your platform. Otherwise the bundled ORT may be missing \
             from your platform package. Try reinstalling decibri.",
            p.display(),
            err
        )),
        None => DecibriError::Other(format!(
            "decibri: failed to initialize ONNX Runtime: {}. \
             Either pass ort_library_path in VadConfig, set ORT_DYLIB_PATH to \
             point to a valid ONNX Runtime library, or enable the \
             `ort-download-binaries` feature for zero-config builds.",
            err
        )),
    }
}

/// Perform the actual ORT init call. Split out by distribution-mode feature
/// so `init_ort_once` stays single-path.
///
/// Under `ort-load-dynamic`: `ort::init_from(path)` is available and is
/// fallible (validates the dylib up-front).
///
/// Under `ort-download-binaries`: `ort::init_from` does NOT exist. ORT is
/// statically linked into the binary and any path argument is meaningless.
/// The path is ignored; we call `ort::init()` to commit an
/// `EnvironmentBuilder` so our OnceLock bookkeeping fires.
#[cfg(all(feature = "vad", feature = "ort-load-dynamic"))]
fn do_ort_init(path: Option<&Path>) -> Result<bool, ort::Error> {
    match path {
        Some(p) => ort::init_from(p).map(|b| b.with_name("decibri").commit()),
        None => Ok(ort::init().with_name("decibri").commit()),
    }
}

#[cfg(all(feature = "vad", not(feature = "ort-load-dynamic")))]
fn do_ort_init(_path: Option<&Path>) -> Result<bool, ort::Error> {
    Ok(ort::init().with_name("decibri").commit())
}

/// Initialize ORT exactly once per process.
///
/// - If ORT is already initialized (by this or any prior caller), returns
///   immediately. The `path` argument is silently ignored. ORT's global
///   state cannot be re-initialized. See `VadConfig::ort_library_path` docs.
/// - Otherwise delegates to the feature-gated `do_ort_init` helper.
/// - On failure, the `OnceLock` is NOT set, so a subsequent caller can retry.
///
/// Note: `e` in the error-wrapping path below is always an `ort::Error` (ORT's
/// own error type), never a `DecibriError`. This function is the single place
/// where ORT errors enter decibri's error hierarchy. Do NOT apply the same
/// wrapping pattern elsewhere or the `decibri:` prefix and guidance string
/// will be duplicated in the message users see.
#[cfg(feature = "vad")]
fn init_ort_once(path: Option<&Path>) -> Result<(), DecibriError> {
    // Fast path: ORT already initialized.
    if ORT_INIT.get().is_some() {
        return Ok(());
    }

    match do_ort_init(path) {
        Ok(_committed) => {
            // First-caller-wins. If another thread set it first, discard ours;
            // ORT's own global init is idempotent (first takes effect).
            let _ = ORT_INIT.set(path.map(|p| p.to_path_buf()));
            Ok(())
        }
        Err(e) => Err(wrap_init_error(path, e)),
    }
}

#[cfg(feature = "vad")]
impl SileroVad {
    /// Create a new Silero VAD instance by loading the ONNX model.
    pub fn new(config: VadConfig) -> Result<Self, DecibriError> {
        let window_size = match config.sample_rate {
            16000 => 512,
            8000 => 256,
            _ => {
                return Err(DecibriError::Other(
                    "Silero VAD only supports sample rates 8000 and 16000".to_string(),
                ))
            }
        };

        if config.threshold < 0.0 || config.threshold > 1.0 {
            return Err(DecibriError::Other(
                "VAD threshold must be between 0.0 and 1.0".to_string(),
            ));
        }

        // Initialize ORT exactly once per process. Subsequent SileroVad
        // instances pass through immediately regardless of their ort_library_path.
        init_ort_once(config.ort_library_path.as_deref())?;

        let session = ort::session::Session::builder()
            .map_err(|e| DecibriError::Other(format!("Failed to create ort session builder: {e}")))?
            .with_intra_threads(1)
            .map_err(|e| DecibriError::Other(format!("Failed to set ort threads: {e}")))?
            .commit_from_file(&config.model_path)
            .map_err(|e| {
                DecibriError::Other(format!(
                    "Failed to load Silero VAD model from {}: {e}",
                    config.model_path.display()
                ))
            })?;

        Ok(Self {
            session,
            state: vec![0.0f32; STATE_SIZE],
            sample_rate: config.sample_rate,
            threshold: config.threshold,
            accumulator: Vec::new(),
            window_size,
        })
    }

    /// Process audio samples and return a VAD result.
    ///
    /// Samples are f32 in the range [-1.0, 1.0]. The chunk can be any size;
    /// internally it is split into `window_size` windows. Leftover samples
    /// carry over to the next call.
    ///
    /// Returns the maximum speech probability across all windows processed.
    /// If no complete windows were formed (chunk too small), returns probability 0.0.
    pub fn process(&mut self, samples: &[f32]) -> Result<VadResult, DecibriError> {
        self.accumulator.extend_from_slice(samples);

        let mut max_probability: f32 = 0.0;
        let mut windows_processed = 0;

        while self.accumulator.len() >= self.window_size {
            let window: Vec<f32> = self.accumulator.drain(..self.window_size).collect();
            let probability = self.infer_window(&window)?;
            max_probability = max_probability.max(probability);
            windows_processed += 1;
        }

        // If no windows processed, return 0 probability (not enough data yet)
        if windows_processed == 0 {
            return Ok(VadResult {
                probability: 0.0,
                is_speech: false,
            });
        }

        Ok(VadResult {
            probability: max_probability,
            is_speech: max_probability >= self.threshold,
        })
    }

    /// Reset LSTM state to zeros (start of new utterance).
    pub fn reset(&mut self) {
        self.state.fill(0.0);
        self.accumulator.clear();
    }

    /// Run inference on a single window of exactly `window_size` samples.
    fn infer_window(&mut self, window: &[f32]) -> Result<f32, DecibriError> {
        // Create input tensors using actual model tensor names:
        //   input: f32[1, window_size]
        //   state: f32[2, 1, 128]
        //   sr:    i64 scalar
        let input_tensor =
            ort::value::Tensor::from_array(([1i64, self.window_size as i64], window.to_vec()))
                .map_err(|e| DecibriError::Other(format!("Failed to create input tensor: {e}")))?;
        let state_tensor =
            ort::value::Tensor::from_array(([2i64, 1i64, 128i64], self.state.clone()))
                .map_err(|e| DecibriError::Other(format!("Failed to create state tensor: {e}")))?;
        let sr_tensor = ort::value::Tensor::from_array(([1i64], vec![self.sample_rate as i64]))
            .map_err(|e| DecibriError::Other(format!("Failed to create sr tensor: {e}")))?;

        let input_values = ort::inputs![
            "input" => input_tensor,
            "state" => state_tensor,
            "sr" => sr_tensor,
        ];

        let outputs = self
            .session
            .run(input_values)
            .map_err(|e| DecibriError::Other(format!("Silero VAD inference failed: {e}")))?;

        // Read outputs using actual model tensor names:
        //   output: f32[1, 1] (speech probability)
        //   stateN: f32[2, 1, 128] (updated state)
        let prob_tensor = outputs["output"]
            .try_extract_tensor::<f32>()
            .map_err(|e| DecibriError::Other(format!("Failed to extract output tensor: {e}")))?;
        let probability = prob_tensor.1[0];

        let state_tensor = outputs["stateN"]
            .try_extract_tensor::<f32>()
            .map_err(|e| DecibriError::Other(format!("Failed to extract state tensor: {e}")))?;
        self.state.copy_from_slice(state_tensor.1);

        Ok(probability)
    }
}

#[cfg(all(test, feature = "vad"))]
mod tests {
    use super::*;
    use std::path::Path;

    fn model_path() -> PathBuf {
        // Resolve relative to the workspace root
        let manifest_dir = env!("CARGO_MANIFEST_DIR");
        Path::new(manifest_dir)
            .join("..")
            .join("..")
            .join("models")
            .join("silero_vad.onnx")
    }

    fn default_config() -> VadConfig {
        VadConfig {
            model_path: model_path(),
            sample_rate: 16000,
            threshold: 0.5,
            ort_library_path: None,
        }
    }

    #[test]
    fn test_wrap_init_error_with_path() {
        use std::env;

        // Platform-agnostic invalid path (guaranteed to not exist on any OS).
        let bogus_path = env::temp_dir().join("does-not-exist-onnxruntime-xyz-test");
        let err = wrap_init_error(Some(bogus_path.as_path()), "simulated ort loader failure");
        let msg = err.to_string();

        assert!(
            msg.contains(&bogus_path.display().to_string()),
            "error message should contain the attempted path, got: {msg}"
        );
        assert!(
            msg.contains("If ORT_DYLIB_PATH is set"),
            "error message should contain actionable guidance phrase, got: {msg}"
        );
        assert!(
            msg.contains("simulated ort loader failure"),
            "error message should include the underlying ort error, got: {msg}"
        );
    }

    #[test]
    fn test_wrap_init_error_without_path() {
        let err = wrap_init_error(None, "simulated ort init failure");
        let msg = err.to_string();

        assert!(
            msg.contains("ort_library_path"),
            "None-path error should mention the VadConfig field, got: {msg}"
        );
        assert!(
            msg.contains("ORT_DYLIB_PATH"),
            "None-path error should mention the env var, got: {msg}"
        );
        assert!(
            msg.contains("ort-download-binaries"),
            "None-path error should mention the opt-out feature, got: {msg}"
        );
        assert!(
            msg.contains("simulated ort init failure"),
            "error message should include the underlying ort error, got: {msg}"
        );
    }

    #[test]
    fn test_vad_config_validation() {
        let bad_rate = VadConfig {
            sample_rate: 44100,
            ..default_config()
        };
        assert!(SileroVad::new(bad_rate).is_err());
    }

    #[test]
    fn test_vad_loads_model() {
        let vad = SileroVad::new(default_config());
        assert!(vad.is_ok(), "Model should load: {:?}", vad.err());
    }

    #[test]
    fn test_vad_silence() {
        let mut vad = SileroVad::new(default_config()).unwrap();
        let silence = vec![0.0f32; 512];
        let result = vad.process(&silence).unwrap();
        // Silence should have low probability
        assert!(
            result.probability < 0.5,
            "Silence probability should be low, got {}",
            result.probability
        );
    }

    #[test]
    fn test_vad_state_persistence() {
        let mut vad = SileroVad::new(default_config()).unwrap();
        let initial_state = vad.state.clone();

        let samples = vec![0.0f32; 512];
        vad.process(&samples).unwrap();

        // State should have changed after inference
        assert_ne!(
            vad.state, initial_state,
            "State should change after inference"
        );
    }

    #[test]
    fn test_vad_accumulator_windows() {
        let mut vad = SileroVad::new(default_config()).unwrap();

        // 1600 samples = 3 windows of 512 + 64 leftover
        let samples = vec![0.0f32; 1600];
        let result = vad.process(&samples).unwrap();
        assert!(result.probability >= 0.0); // Just verify it runs

        // 64 samples should remain in accumulator
        assert_eq!(vad.accumulator.len(), 64);
    }

    #[test]
    fn test_vad_accumulator_carry() {
        let mut vad = SileroVad::new(default_config()).unwrap();

        // First chunk: 1600 samples → 3 windows, 64 leftover
        let chunk1 = vec![0.0f32; 1600];
        vad.process(&chunk1).unwrap();
        assert_eq!(vad.accumulator.len(), 64);

        // Second chunk: 1600 samples → 64 + 1600 = 1664 → 3 windows, 128 leftover
        let chunk2 = vec![0.0f32; 1600];
        vad.process(&chunk2).unwrap();
        assert_eq!(vad.accumulator.len(), 128);
    }

    #[test]
    fn test_vad_small_chunk() {
        let mut vad = SileroVad::new(default_config()).unwrap();

        // Less than one window, so no inference should run
        let samples = vec![0.0f32; 100];
        let result = vad.process(&samples).unwrap();
        assert_eq!(result.probability, 0.0);
        assert_eq!(vad.accumulator.len(), 100);
    }

    #[test]
    fn test_vad_reset() {
        let mut vad = SileroVad::new(default_config()).unwrap();

        // Process something to change state
        let samples = vec![0.0f32; 512];
        vad.process(&samples).unwrap();
        vad.accumulator.extend_from_slice(&[0.0; 100]); // add some leftover

        vad.reset();

        assert!(vad.state.iter().all(|&v| v == 0.0), "State should be zeros");
        assert!(vad.accumulator.is_empty(), "Accumulator should be empty");
    }
}