path: root/crates/asist-audio/src/vad.rs

//! Detección de voz y troceado en intervenciones.
//!
//! Está escrito como una máquina de estados pura: se le dan muestras y
//! devuelve eventos, sin hilos ni canales dentro. Así el comportamiento que
//! más cuesta depurar a oído —cuándo arranca un turno, cuándo lo corta el
//! silencio, cuándo se ignora el eco del propio altavoz— se puede probar
//! entero con audio sintético y sin micrófono.

use std::collections::VecDeque;

use asist_core::config::VadConfig;

use crate::{rms, ASR_SAMPLE_RATE};

/// Una intervención cerrada, lista para transcribir.
#[derive(Debug, Clone)]
pub struct Utterance {
    pub samples: Vec<f32>,
    pub sample_rate: u32,
}

impl Utterance {
    pub fn duration_secs(&self) -> f32 {
        self.samples.len() as f32 / self.sample_rate as f32
    }
}

/// Lo que el segmentador tiene que contar hacia fuera.
#[derive(Debug, Clone)]
pub enum VoiceEvent {
    /// Ha empezado a hablarse.
    Started,
    /// Audio nuevo dentro de la intervención en curso, para las
    /// transcripciones provisionales.
    Audio(Vec<f32>),
    /// Intervención terminada y lo bastante larga como para transcribirla.
    Ended(Utterance),
    /// Terminada pero demasiado corta: un golpe en la mesa, una tos.
    Discarded,
    /// Se ha detectado voz mientras el asistente hablaba, con el barge-in
    /// activo. El orquestador corta la reproducción al recibirlo.
    BargeIn,
}

/// Qué hace el segmentador con el micrófono mientras suena el altavoz.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Gate {
    /// Nadie está hablando por el altavoz: se escucha con normalidad.
    Open,
    /// El asistente habla. Según la configuración, o se ignora la entrada
    /// (media dúplex) o se exige más volumen para interrumpir (barge-in).
    Speaking,
}

pub struct Segmenter {
    config: VadConfig,
    frame_samples: usize,
    preroll_samples: usize,
    silence_hold_samples: usize,
    min_utterance_samples: usize,
    max_utterance_samples: usize,

    pending: Vec<f32>,
    preroll: VecDeque<f32>,
    utterance: Vec<f32>,
    /// Muestras de la intervención que estaban de verdad por encima del
    /// umbral. El mínimo se mide sobre esto y no sobre `utterance`, que
    /// arrastra el preroll: si no, 0,1 s de golpe en la mesa más 0,2 s de
    /// preroll pasan por una intervención válida y disparan un turno entero.
    voiced: usize,
    noise_floor: f32,
    silence_run: usize,
    speaking: bool,
    gate: Gate,
}

impl Segmenter {
    pub fn new(config: &VadConfig) -> Self {
        let rate = ASR_SAMPLE_RATE as f32;
        Self {
            frame_samples: (rate * config.frame_seconds).max(1.0) as usize,
            preroll_samples: (rate * config.preroll_seconds) as usize,
            silence_hold_samples: (rate * config.silence_hold) as usize,
            min_utterance_samples: (rate * config.min_utterance) as usize,
            max_utterance_samples: (rate * config.max_utterance) as usize,
            config: config.clone(),
            pending: Vec::new(),
            preroll: VecDeque::new(),
            utterance: Vec::new(),
            voiced: 0,
            noise_floor: 0.0,
            silence_run: 0,
            speaking: false,
            gate: Gate::Open,
        }
    }

    /// Abre o cierra el micrófono según hable o no el asistente.
    pub fn set_gate(&mut self, gate: Gate) {
        if self.gate == gate {
            return;
        }
        self.gate = gate;
        // Al volver a abrir tras una respuesta, lo acumulado es la cola del
        // propio altavoz: arrancar un turno con eso daría un turno fantasma.
        if gate == Gate::Open {
            self.pending.clear();
            self.preroll.clear();
            self.utterance.clear();
            self.voiced = 0;
            self.silence_run = 0;
            self.speaking = false;
        }
    }

    pub fn is_speaking(&self) -> bool {
        self.speaking
    }

    pub fn noise_floor(&self) -> f32 {
        self.noise_floor
    }

    /// Umbral que separa voz de silencio ahora mismo.
    pub fn threshold(&self) -> f32 {
        let base = (self.noise_floor * self.config.threshold_factor)
            .clamp(self.config.min_threshold, self.config.max_threshold);
        // Con el altavoz sonando hay que hablar más alto para colarse: el
        // micrófono se está oyendo a sí mismo.
        if self.gate == Gate::Speaking {
            base * self.config.barge_in_factor
        } else {
            base
        }
    }

    /// Cierra a la fuerza la intervención en curso (cierre del programa).
    pub fn flush(&mut self) -> Option<Utterance> {
        if !self.speaking || self.voiced < self.min_utterance_samples {
            return None;
        }
        self.speaking = false;
        self.voiced = 0;
        Some(Utterance {
            samples: std::mem::take(&mut self.utterance),
            sample_rate: ASR_SAMPLE_RATE,
        })
    }

    /// Alimenta audio mono a 16 kHz y recoge lo que haya que hacer.
    pub fn push(&mut self, samples: &[f32]) -> Vec<VoiceEvent> {
        let mut events = Vec::new();
        // En media dúplex el micrófono está apagado de hecho: sin esto, el
        // asistente se transcribe a sí mismo y se responde solo.
        if self.gate == Gate::Speaking && !self.config.barge_in {
            return events;
        }
        self.pending.extend_from_slice(samples);

        while self.pending.len() >= self.frame_samples {
            let frame: Vec<f32> = self.pending.drain(..self.frame_samples).collect();
            let level = rms(&frame);
            self.track_noise_floor(level);
            let threshold = self.threshold();

            if level < threshold && !self.speaking {
                self.preroll.extend(frame.iter().copied());
                while self.preroll.len() > self.preroll_samples {
                    self.preroll.pop_front();
                }
                continue;
            }

            if !self.speaking {
                self.speaking = true;
                self.utterance.clear();
                self.voiced = 0;
                self.utterance.extend(self.preroll.drain(..));
                if self.gate == Gate::Speaking {
                    events.push(VoiceEvent::BargeIn);
                }
                events.push(VoiceEvent::Started);
            }

            if level < threshold {
                self.silence_run += frame.len();
            } else {
                self.silence_run = 0;
                self.voiced += frame.len();
            }
            self.utterance.extend_from_slice(&frame);

            let ended = self.silence_run >= self.silence_hold_samples;
            let too_long = self.utterance.len() >= self.max_utterance_samples;
            if !ended && !too_long {
                events.push(VoiceEvent::Audio(frame));
                continue;
            }

            let long_enough = self.voiced >= self.min_utterance_samples;
            events.push(if long_enough {
                VoiceEvent::Ended(Utterance {
                    samples: std::mem::take(&mut self.utterance),
                    sample_rate: ASR_SAMPLE_RATE,
                })
            } else {
                VoiceEvent::Discarded
            });

            self.utterance.clear();
            self.voiced = 0;
            self.preroll.clear();
            self.silence_run = 0;
            // Un corte por longitud cae a mitad de frase: se sigue escuchando
            // como si el usuario no hubiera dejado de hablar, que es la verdad.
            self.speaking = too_long && !ended;
            if self.speaking {
                events.push(VoiceEvent::Started);
            }
        }
        events
    }

    /// El suelo de ruido sólo baja: una voz sostenida no debe poder arrastrar
    /// el umbral por encima de sí misma y dejar de detectarse.
    fn track_noise_floor(&mut self, level: f32) {
        if self.noise_floor == 0.0 {
            self.noise_floor = level;
        } else if level < self.noise_floor * 1.5 {
            self.noise_floor = self.noise_floor * 0.95 + level * 0.05;
        }
    }
}

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

    fn config() -> VadConfig {
        VadConfig {
            frame_seconds: 0.1,
            preroll_seconds: 0.2,
            silence_hold: 0.3,
            min_utterance: 0.3,
            max_utterance: 2.0,
            threshold_factor: 3.0,
            min_threshold: 0.001,
            max_threshold: 0.02,
            barge_in: false,
            barge_in_factor: 4.0,
        }
    }

    fn samples(secs: f32, amplitude: f32) -> Vec<f32> {
        let n = (ASR_SAMPLE_RATE as f32 * secs) as usize;
        // Alterna de signo para que el RMS sea la amplitud y no un continuo.
        (0..n)
            .map(|i| if i % 2 == 0 { amplitude } else { -a