//! 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 { -amplitude })
            .collect()
    }

    fn feed(seg: &mut Segmenter, secs: f32, amplitude: f32) -> Vec<VoiceEvent> {
        seg.push(&samples(secs, amplitude))
    }

    #[test]
    fn el_silencio_no_arranca_ningun_turno() {
        let mut seg = Segmenter::new(&config());
        let events = feed(&mut seg, 2.0, 0.0001);
        assert!(
            events.is_empty(),
            "el silencio no debe producir eventos: {events:?}"
        );
    }

    #[test]
    fn la_voz_seguida_de_silencio_cierra_una_intervencion() {
        let mut seg = Segmenter::new(&config());
        feed(&mut seg, 1.0, 0.0002); // deja que el suelo de ruido se asiente
        let mut events = feed(&mut seg, 0.8, 0.3);
        events.extend(feed(&mut seg, 0.6, 0.0002));

        assert!(matches!(events.first(), Some(VoiceEvent::Started)));
        let ended = events.iter().find_map(|e| match e {
            VoiceEvent::Ended(u) => Some(u),
            _ => None,
        });
        let utterance = ended.expect("la intervención debió cerrarse");
        assert!(
            utterance.duration_secs() > 0.8,
            "el preroll debe ir incluido, duró {}",
            utterance.duration_secs()
        );
    }

    #[test]
    fn un_ruido_corto_se_descarta() {
        let mut seg = Segmenter::new(&config());
        feed(&mut seg, 1.0, 0.0002);
        let mut events = feed(&mut seg, 0.1, 0.3);
        events.extend(feed(&mut seg, 0.6, 0.0002));
        // 0,1 s de golpe no llegan al mínimo de 0,3 s de voz, por mucho que el
        // preroll haga que la intervención dure más.
        assert!(
            events.iter().any(|e| matches!(e, VoiceEvent::Discarded)),
            "esperaba un descarte: {events:?}"
        );
    }

    #[test]
    fn una_intervencion_interminable_se_corta_y_se_sigue_escuchando() {
        let mut seg = Segmenter::new(&config());
        feed(&mut seg, 1.0, 0.0002);
        let events = feed(&mut seg, 3.0, 0.3);
        assert!(
            events.iter().any(|e| matches!(e, VoiceEvent::Ended(_))),
            "a los 2 s debe cortarse: {events:?}"
        );
        assert!(
            seg.is_speaking(),
            "tras un corte forzado se sigue en mitad de la frase"
        );
    }

    #[test]
    fn en_media_duplex_el_altavoz_no_se_transcribe_a_si_mismo() {
        let mut seg = Segmenter::new(&config());
        feed(&mut seg, 1.0, 0.0002);
        seg.set_gate(Gate::Speaking);
        let events = feed(&mut seg, 2.0, 0.5);
        assert!(
            events.is_empty(),
            "con barge_in apagado no debe entrar nada mientras habla el asistente: {events:?}"
        );
    }

    #[test]
    fn con_barge_in_hace_falta_hablar_mas_alto_para_cortar() {
        let mut config = config();
        config.barge_in = true;
        config.barge_in_factor = 4.0;
        let mut seg = Segmenter::new(&config);
        feed(&mut seg, 1.0, 0.0002);
        seg.set_gate(Gate::Speaking);

        // Justo por encima del umbral normal pero por debajo del elevado.
        let eco = seg.threshold() / config.barge_in_factor * 1.5;
        let events = feed(&mut seg, 0.5, eco);
        assert!(
            events.is_empty(),
            "el eco del altavoz no debe cortar: {events:?}"
        );

        let voz = seg.threshold() * 2.0;
        let events = feed(&mut seg, 0.5, voz);
        assert!(
            events.iter().any(|e| matches!(e, VoiceEvent::BargeIn)),
            "una voz clara sí debe cortar: {events:?}"
        );
    }

    #[test]
    fn al_reabrir_el_microfono_se_tira_la_cola_del_altavoz() {
        let mut config = config();
        config.barge_in = true;
        let mut seg = Segmenter::new(&config);
        feed(&mut seg, 1.0, 0.0002);
        seg.set_gate(Gate::Speaking);
        feed(&mut seg, 0.5, 0.9);
        assert!(seg.is_speaking());

        seg.set_gate(Gate::Open);
        assert!(
            !seg.is_speaking(),
            "reabrir debe descartar lo acumulado; si no, el turno siguiente arranca con eco"
        );
    }

    #[test]
    fn el_suelo_de_ruido_no_sube_con_la_voz() {
        let mut seg = Segmenter::new(&config());
        feed(&mut seg, 1.0, 0.0002);
        let quieto = seg.noise_floor();
        feed(&mut seg, 2.0, 0.5);
        assert!(
            seg.noise_floor() <= quieto * 1.5,
            "hablar no debe elevar el suelo de ruido ({quieto} -> {})",
            seg.noise_floor()
        );
    }

    #[test]
    fn el_cierre_entrega_la_intervencion_a_medias() {
        let mut seg = Segmenter::new(&config());
        feed(&mut seg, 1.0, 0.0002);
        feed(&mut seg, 0.5, 0.3);
        assert!(
            seg.flush().is_some(),
            "lo ya hablado no debe perderse al cerrar"
        );
        assert!(seg.flush().is_none(), "y no debe entregarse dos veces");
    }
}