//! Voice activity detection and splitting into utterances.
//!
//! It is written as a pure state machine: it is fed samples and returns
//! events, with no threads or channels inside. That way the behaviour that is
//! hardest to debug by ear (when a turn starts, when silence ends it, when
//! the echo of the assistant's own speaker is ignored) can be fully tested
//! with synthetic audio and no microphone.

use std::collections::VecDeque;

use asist_core::config::VadConfig;

use crate::{rms, ASR_SAMPLE_RATE};

/// A closed utterance, ready to be transcribed.
#[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
    }
}

/// What the segmenter reports to the outside.
#[derive(Debug, Clone)]
pub enum VoiceEvent {
    /// Ha empezado a hablarse.
    Started,
    /// New audio inside the current utterance, for the partial
    /// transcriptions.
    Audio(Vec<f32>),
    /// Utterance finished and long enough to be transcribed.
    Ended(Utterance),
    /// Finished but too short: a knock on the table, a cough.
    Discarded,
    /// Speech was detected while the assistant was talking, with barge-in
    /// on. The orchestrator cuts playback when it gets this.
    BargeIn,
}

/// What the segmenter does with the microphone while the speaker plays.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Gate {
    /// Nothing is playing on the speaker: listen normally.
    Open,
    /// The assistant is talking. Depending on the configuration, input is
    /// either ignored (half duplex) or needs more volume to interrupt (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>,
    /// Samples of the utterance that really were above the threshold. The
    /// minimum is measured on this and not on `utterance`, which carries the
    /// preroll: otherwise 0.1 s of a knock on the table plus 0.2 s of preroll
    /// passes for a valid utterance and triggers a whole turn.
    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,
        }
    }

    /// Opens or closes the microphone depending on whether the assistant talks.
    pub fn set_gate(&mut self, gate: Gate) {
        if self.gate == gate {
            return;
        }
        self.gate = gate;
        // When reopening after an answer, what has accumulated is the tail of
        // the assistant's own speaker: starting a turn with it would be a ghost turn.
        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
    }

    /// Threshold that separates speech from silence right now.
    pub fn threshold(&self) -> f32 {
        let base = (self.noise_floor * self.config.threshold_factor)
            .clamp(self.config.min_threshold, self.config.max_threshold);
        // With the speaker playing you must speak louder to get through: the
        // microphone is hearing itself.
        if self.gate == Gate::Speaking {
            base * self.config.barge_in_factor
        } else {
            base
        }
    }

    /// Forcibly closes the current utterance (program shutdown).
    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,
        })
    }

    /// Feeds 16 kHz mono audio and collects whatever needs doing.
    pub fn push(&mut self, samples: &[f32]) -> Vec<VoiceEvent> {
        let mut events = Vec::new();
        // In half duplex the microphone is effectively off: without this, the
        // assistant transcribes itself and answers itself.
        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;
            // A length cut lands mid-sentence: keep listening as if the user had
            // not stopped talking, which is the truth.
            self.speaking = too_long && !ended;
            if self.speaking {
                events.push(VoiceEvent::Started);
            }
        }
        events
    }

    /// The noise floor only goes down: sustained speech must not be able to
    /// drag the threshold above itself and stop being detected.
    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;
        // Alternates sign so the RMS is the amplitude and not a DC level.
        (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 silence_starts_no_turn() {
        let mut seg = Segmenter::new(&config());
        let events = feed(&mut seg, 2.0, 0.0001);
        assert!(
            events.is_empty(),
            "silence must produce no events: {events:?}"
        );
    }

    #[test]
    fn speech_followed_by_silence_closes_an_utterance() {
        let mut seg = Segmenter::new(&config());
        feed(&mut seg, 1.0, 0.0002); // let the noise floor settle
        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("the utterance should have closed");
        assert!(
            utterance.duration_secs() > 0.8,
            "the preroll must be included, it lasted {}",
            utterance.duration_secs()
        );
    }

    #[test]
    fn short_noise_is_discarded() {
        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 of a knock does not reach the 0.3 s speech minimum, however
        // much the preroll makes the utterance last longer.
        assert!(
            events.iter().any(|e| matches!(e, VoiceEvent::Discarded)),
            "expected a discard: {events:?}"
        );
    }

    #[test]
    fn endless_utterance_is_cut_and_listening_continues() {
        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(_))),
            "it must be cut at 2 s: {events:?}"
        );
        assert!(
            seg.is_speaking(),
            "after a forced cut we are still mid-sentence"
        );
    }

    #[test]
    fn in_half_duplex_the_speaker_is_not_transcribed() {
        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(),
            "with barge_in off nothing may come in while the assistant talks: {events:?}"
        );
    }

    #[test]
    fn with_barge_in_you_must_speak_louder_to_interrupt() {
        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);

        // Just above the normal threshold but below the raised one.
        let eco = seg.threshold() / config.barge_in_factor * 1.5;
        let events = feed(&mut seg, 0.5, eco);
        assert!(
            events.is_empty(),
            "speaker echo must not interrupt: {events:?}"
        );

        let speech = seg.threshold() * 2.0;
        let events = feed(&mut seg, 0.5, speech);
        assert!(
            events.iter().any(|e| matches!(e, VoiceEvent::BargeIn)),
            "clear speech must interrupt: {events:?}"
        );
    }

    #[test]
    fn reopening_the_microphone_drops_the_speaker_tail() {
        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(),
            "reopening must drop what accumulated; otherwise the next turn starts with echo"
        );
    }

    #[test]
    fn noise_floor_does_not_rise_with_speech() {
        let mut seg = Segmenter::new(&config());
        feed(&mut seg, 1.0, 0.0002);
        let quiet = seg.noise_floor();
        feed(&mut seg, 2.0, 0.5);
        assert!(
            seg.noise_floor() <= quiet * 1.5,
            "speech must not raise the noise floor ({quiet} -> {})",
            seg.noise_floor()
        );
    }

    #[test]
    fn shutdown_delivers_the_unfinished_utterance() {
        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");
    }
}