piper-phoneme-streaming 0.1.1

use std::collections::VecDeque;

use crate::semantic::Language;
use crate::text_expand::ExpandUnit;

const CONTEXT_WINDOW_SIZE: usize = 5;
/// A single unambiguous word with confidence ≥ this triggers an immediate
/// language switch and resets the context window (fast path).
const SINGLE_WORD_SWITCH_THRESHOLD: f64 = 0.80;
/// The context window must reach this confidence before switching language
/// (slow path).  0.70 = 0.5 base + 0.20 hysteresis.
const CONTEXT_SWITCH_THRESHOLD: f64 = 0.70;

/// Trait for language detection backends.
///
/// Implementations receive a context string and return the most likely language
/// plus its confidence score, or `None` if detection fails.
pub(crate) trait LanguageDetector: Send + Sync {
    fn detect(&self, context: &str) -> Option<(Language, f64)>;
}

/// Streaming language detector that wraps a [`LanguageDetector`] backend.
///
/// Takes [`ExpandUnit`] values one at a time and returns the detected
/// [`Language`] for each:
///
/// - `Word` units trigger detection via a sliding context window with
///   hysteresis — language only switches when confidence exceeds 0.5 + 0.20.
/// - `Number` and `Mark` units inherit the current language without detection.
/// - Sentence boundaries (`.`, `?`, `!`) should be signalled via
///   [`reset_context`](Self::reset_context) to clear the sliding window.
pub struct StreamingLanguageDetector {
    detector: Box<dyn LanguageDetector>,
    current_language: Language,
    context_window: VecDeque<String>,
}

impl StreamingLanguageDetector {
    pub(crate) fn new(default_language: Language, detector: Box<dyn LanguageDetector>) -> Self {
        Self {
            detector,
            current_language: default_language,
            context_window: VecDeque::new(),
        }
    }

    /// Build a detector backed by the `lingua` library.
    ///
    /// `languages` must contain at least two languages; `default_language` is
    /// the starting assumption before any text is seen.
    ///
    pub fn with_lingua(languages: &[Language], default_language: Language) -> Self {
        Self::new(default_language, Box::new(LinguaDetector::new(languages)))
    }

    /// Push one expand unit and get back the language for that unit.
    ///
    /// Words run through the detection algorithm; numbers and marks inherit
    /// the current language without running detection.
    pub fn push(&mut self, unit: &ExpandUnit) -> Language {
        match unit {
            ExpandUnit::Word(word) => self.detect_for_word(word),
            ExpandUnit::Number(_) | ExpandUnit::Mark(_) => self.current_language,
        }
    }

    /// Clear the context window on sentence boundaries (`.`, `?`, `!`).
    pub fn reset_context(&mut self) {
        self.context_window.clear();
    }

    fn detect_for_word(&mut self, word: &str) -> Language {
        // Fast path: a single unambiguous word with high confidence triggers an
        // immediate switch and resets the context window so the new language is
        // not dragged back by old context.
        if let Some((word_lang, word_conf)) = self.detector.detect(word)
            && word_lang != self.current_language
            && word_conf >= SINGLE_WORD_SWITCH_THRESHOLD
        {
            self.current_language = word_lang;
            self.context_window.clear();
            self.context_window.push_back(word.to_string());
            return self.current_language;
        }

        // Slow path: accumulate a sliding context window and switch only when
        // the aggregated confidence exceeds the hysteresis threshold.
        self.context_window.push_back(word.to_string());
        if self.context_window.len() > CONTEXT_WINDOW_SIZE {
            self.context_window.pop_front();
        }

        let context: String = self
            .context_window
            .iter()
            .map(String::as_str)
            .collect::<Vec<_>>()
            .join(" ");

        if let Some((ctx_lang, ctx_conf)) = self.detector.detect(&context)
            && ctx_lang != self.current_language
            && ctx_conf >= CONTEXT_SWITCH_THRESHOLD
        {
            self.current_language = ctx_lang;
        }

        self.current_language
    }
}

// ---------------------------------------------------------------------------
// lingua backend
// ---------------------------------------------------------------------------

struct LinguaDetector {
    detector: lingua::LanguageDetector,
}

impl LinguaDetector {
    fn new(languages: &[Language]) -> Self {
        let lingua_langs: Vec<lingua::Language> = languages
            .iter()
            .map(|l| match l {
                Language::English => lingua::Language::English,
                Language::Vietnamese => lingua::Language::Vietnamese,
            })
            .collect();
        let detector = lingua::LanguageDetectorBuilder::from_languages(&lingua_langs)
            .with_minimum_relative_distance(0.25)
            .build();
        Self { detector }
    }
}

impl LanguageDetector for LinguaDetector {
    fn detect(&self, context: &str) -> Option<(Language, f64)> {
        let confidences = self.detector.compute_language_confidence_values(context);
        confidences.first().map(|(lingua_lang, confidence)| {
            let lang = match lingua_lang {
                lingua::Language::English => Language::English,
                lingua::Language::Vietnamese => Language::Vietnamese,
            };
            (lang, *confidence)
        })
    }
}

// ---------------------------------------------------------------------------
// Tests
//
// Each test case is: (preferred_language, &[(word, expected_language)])
// Every word is pushed as ExpandUnit::Word; the output language is asserted
// per word so failures point to the exact token where behavior diverges.
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::text_expand::ExpandUnit;
    use Language::{English as EN, Vietnamese as VI};

    fn det(default: Language) -> StreamingLanguageDetector {
        StreamingLanguageDetector::with_lingua(&[EN, VI], default)
    }

    /// Core helper: push each `(word, expected_language)` pair and assert per step.
    fn check(preferred: Language, steps: &[(&str, Language)]) {
        let mut d = det(preferred);
        for (word, expected) in steps {
            let got = d.push(&ExpandUnit::Word(word.to_string()));
            assert_eq!(
                got, *expected,
                "preferred={preferred:?}  word={word:?}  expected={expected:?}  got={got:?}"
            );
        }
    }

    // -------------------------------------------------------------------------
    // Pure-language — window fills with one language, never switches
    // -------------------------------------------------------------------------

    #[test]
    fn case_pure_english() {
        check(
            EN,
            &[
                ("the", EN),   // window: [the]
                ("quick", EN), // window: [the, quick]
                ("brown", EN), // window: [the, quick, brown]
                ("fox", EN),   // window: [the, quick, brown, fox]
                ("jumps", EN), // window: [the, quick, brown, fox, jumps]
            ],
        );
    }

    #[test]
    fn case_pure_vietnamese() {
        check(
            VI,
            &[
                ("xin", VI),  // window: [xin]
                ("chào", VI), // window: [xin, chào]
                ("bạn", VI),  // window: [xin, chào, bạn]
                ("tên", VI),  // window: [xin, chào, bạn, tên]
                ("là", VI),   // window: [xin, chào, bạn, tên, là]
            ],
        );
    }

    // -------------------------------------------------------------------------
    // EN → VI transition
    //
    // "chào" has single-word VI confidence = 1.0 ≥ 0.80 → fast switch at
    // first clear Vietnamese word.  Context resets; VI stable thereafter.
    // -------------------------------------------------------------------------

    #[test]
    fn case_en_to_vi_transition() {
        check(
            EN,
            &[
                // ── build EN context ──
                ("the", EN),   // window: [the]
                ("quick", EN), // window: [the, quick]
                ("brown", EN), // window: [the, quick, brown]
                ("fox", EN),   // window: [the, quick, brown, fox]
                ("jumps", EN), // window: [the, quick, brown, fox, jumps]
                // ── first VI word: high single-word confidence → fast switch ──
                ("chào", VI), // fast path: VI 1.0 ≥ 0.80 → switch; context reset to [chào]
                // ── settled in VI ──
                ("bạn", VI),
                ("tôi", VI),
                ("muốn", VI),
                ("học", VI),
            ],
        );
    }

    // -------------------------------------------------------------------------
    // VI → EN transition
    //
    // "the" has single-word EN confidence ≥ 0.80 → fast switch immediately.
    // Context resets to ["the"]; EN stable for all following EN words.
    // -------------------------------------------------------------------------

    #[test]
    fn case_vi_to_en_transition() {
        check(
            VI,
            &[
                // ── build VI context ──
                ("xin", VI),
                ("chào", VI),
                ("bạn", VI),
                ("tên", VI),
                ("là", VI),
                // ── first EN word: fast switch ──
                ("the", EN), // fast path: EN ≥ 0.80 → switch; context reset to [the]
                // ── settled in EN ──
                ("quick", EN),
                ("brown", EN),
                ("fox", EN),
                ("jumps", EN),
            ],
        );
    }

    // -------------------------------------------------------------------------
    // Reset clears context — next language dominates immediately
    //
    // With an empty window after reset, even a single unambiguous word exceeds
    // the 0.70 threshold and switches language.
    // -------------------------------------------------------------------------

    #[test]
    fn case_reset_vi_context_then_en() {
        let mut d = det(VI);
        for w in &["xin", "chào", "bạn"] {
            d.push(&ExpandUnit::Word(w.to_string()));
        }
        d.reset_context();
        // Empty window → single EN word immediately switches
        for (word, expected) in &[
            ("the", EN), // window: [the]   — triggers switch on first word
            ("quick", EN),
            ("brown", EN),
            ("fox", EN),
            ("jumps", EN),
        ] {
            let got = d.push(&ExpandUnit::Word(word.to_string()));
            assert_eq!(
                got, *expected,
                "after reset  word={word:?}  expected={expected:?}  got={got:?}"
            );
        }
    }

    // -------------------------------------------------------------------------
    // Numbers and marks — always inherit current language, never enter window
    // -------------------------------------------------------------------------

    #[test]
    fn case_number_inherits_en() {
        let mut d = det(EN);
        assert_eq!(d.push(&ExpandUnit::Number("42".into())), EN);
        assert_eq!(d.push(&ExpandUnit::Number("0".into())), EN);
    }

    #[test]
    fn case_number_inherits_vi() {
        let mut d = det(VI);
        assert_eq!(d.push(&ExpandUnit::Number("100".into())), VI);
        assert_eq!(d.push(&ExpandUnit::Number("1000".into())), VI);
    }

    #[test]
    fn case_mark_inherits_en() {
        let mut d = det(EN);
        assert_eq!(d.push(&ExpandUnit::Mark(' ')), EN);
        assert_eq!(d.push(&ExpandUnit::Mark(',')), EN);
        assert_eq!(d.push(&ExpandUnit::Mark('.')), EN);
    }

    #[test]
    fn case_mark_inherits_vi() {
        let mut d = det(VI);
        assert_eq!(d.push(&ExpandUnit::Mark(' ')), VI);
        assert_eq!(d.push(&ExpandUnit::Mark(',')), VI);
    }

    /// Numbers do not enter the context window; they inherit the current language.
    /// With fast-path switching, "chào" switches to VI immediately.
    #[test]
    fn case_number_does_not_affect_transition_timing() {
        let mut d = det(EN);
        for w in &["the", "quick", "brown", "fox", "jumps"] {
            d.push(&ExpandUnit::Word(w.to_string()));
        }
        assert_eq!(d.push(&ExpandUnit::Word("chào".into())), VI); // fast path: VI 1.0 → switch
        assert_eq!(d.push(&ExpandUnit::Number("100".into())), VI); // inherits VI
        assert_eq!(d.push(&ExpandUnit::Word("bạn".into())), VI); // still VI
        assert_eq!(d.push(&ExpandUnit::Number("7".into())), VI); // inherits VI
    }

    /// Marks behave identically — no window effect.
    #[test]
    fn case_mark_does_not_affect_transition_timing() {
        let mut d = det(EN);
        for w in &["the", "quick", "brown", "fox", "jumps"] {
            d.push(&ExpandUnit::Word(w.to_string()));
        }
        assert_eq!(d.push(&ExpandUnit::Word("chào".into())), VI); // fast path: VI 1.0 → switch
        assert_eq!(d.push(&ExpandUnit::Mark(',')), VI); // inherits VI
        assert_eq!(d.push(&ExpandUnit::Word("bạn".into())), VI); // still VI
    }

    // -------------------------------------------------------------------------
    // Mixed sentences — fast path switches on clear single-word signals
    //
    // "trong tiếng anh hello world có nghĩa là xin chào"
    // = "in English, 'hello world' means 'xin chào'"
    //
    // "hello" (EN 0.886) and "world" (EN 0.851) both exceed the fast-path
    // threshold → immediate EN switch.  "có" exceeds the VI threshold →
    // fast switch back to VI.
    // -------------------------------------------------------------------------

    #[test]
    fn case_vi_sentence_with_embedded_en_words() {
        check(
            VI,
            &[
                // ── pure VI context ──
                ("trong", VI),
                ("tiếng", VI),
                ("anh", VI),
                // ── embedded English words — fast-path switches to EN ──
                ("hello", EN), // fast path: EN 0.886 ≥ 0.80 → switch; context reset to [hello]
                ("world", EN), // stays EN (EN 0.851 ≥ 0.80 but already EN)
                // ── back to Vietnamese — fast-path switches back ──
                ("có", VI), // fast path: VI ≥ 0.80 → switch back
                ("nghĩa", VI),
                ("là", VI),
                ("xin", VI),
                ("chào", VI),
            ],
        );
    }

    /// A single Vietnamese word with high single-word confidence fast-switches
    /// out of EN, and the next clear English word fast-switches back.
    #[test]
    fn case_en_sentence_with_single_embedded_vi_word() {
        check(
            EN,
            &[
                ("the", EN),
                ("quick", EN),
                ("brown", EN),
                ("fox", EN),
                // ── single VI word: VI 1.0 ≥ 0.80 → fast switch ──
                ("chào", VI), // fast path: VI 1.0 → switch; context reset to [chào]
                // ── next clear EN word fast-switches back ──
                ("jumps", EN), // fast path: EN ≥ 0.80 → switch back
                ("over", EN),
                ("lazy", EN),
                ("dog", EN),
                ("today", EN),
            ],
        );
    }

    // -------------------------------------------------------------------------
    // Numbers in mixed-language context
    // -------------------------------------------------------------------------

    /// "giá 100 đồng" — number throughout inherits VI; reset then switches EN.
    #[test]
    fn case_number_in_vi_then_reset_to_en() {
        let mut d = det(VI);
        for w in &["giá", "tiền", "là"] {
            d.push(&ExpandUnit::Word(w.to_string()));
        }
        assert_eq!(d.push(&ExpandUnit::Number("100".into())), VI);
        d.reset_context();
        assert_eq!(d.push(&ExpandUnit::Word("the".into())), EN); // immediate switch after reset
        assert_eq!(d.push(&ExpandUnit::Word("price".into())), EN);
    }

    /// Number during the EN→VI transition: with fast-path, "chào" switches
    /// immediately and the number inherits VI right away.
    #[test]
    fn case_number_tracks_language_through_en_to_vi_transition() {
        let mut d = det(EN);
        for w in &["the", "quick", "brown", "fox", "jumps"] {
            d.push(&ExpandUnit::Word(w.to_string()));
        }
        assert_eq!(d.push(&ExpandUnit::Word("chào".into())), VI); // fast path: VI 1.0 → switch
        assert_eq!(d.push(&ExpandUnit::Number("42".into())), VI); // inherits VI immediately
        assert_eq!(d.push(&ExpandUnit::Word("bạn".into())), VI); // still VI
        assert_eq!(d.push(&ExpandUnit::Number("7".into())), VI); // inherits VI
    }
}