piper-phoneme-streaming 0.1.1

use crate::G2pToken;
use crate::compat_espeak::translate::ipa_table::PendingStress;
use crate::compat_espeak::translate::phonemes_to_ipa_full;
use crate::dictionary::stress::{change_word_stress, promote_strend_stress};
use crate::embedded_data::materialized_data_dir;
use crate::error::Result;
use crate::phoneme::PhonemeData;
use crate::semantic::{Language, SentenceUnit, WordPhoneme};
use std::sync::Arc;

mod full;
mod streaming;

pub use self::full::FullSentencePhonemeUpgrade;
pub use self::streaming::{
    StreamingSentencePhonemeUpgrade, StreamingSentencePhonemeUpgradeSession,
};

const FLAG_STREND: u32 = 1 << 9;
const FLAG_STREND2: u32 = 1 << 10;
const PRIMARY_A: u8 = 6;
const PRIMARY_B: u8 = 7;

#[derive(Clone)]
pub(crate) struct Entry {
    pub(crate) kind: Kind,
    pub(crate) raw_phonemes: Vec<u8>,
    pub(crate) flags: u32,
    pub(crate) normalized_word: Option<String>,
    pub(crate) passthrough_unit: Option<SentenceUnit>,
    pub(crate) language: Option<Language>,
}

#[derive(Clone, Copy, PartialEq, Eq)]
pub(crate) enum Kind {
    Word,
    ClauseBoundary,
    Other,
}

impl Entry {
    pub(crate) fn from_sentence_unit(unit: &SentenceUnit) -> Self {
        match unit {
            SentenceUnit::Word(word) => Self {
                kind: Kind::Word,
                raw_phonemes: word.raw_codes().to_vec(),
                flags: word.flags.raw(),
                normalized_word: Some(word.normalized_word.clone()),
                passthrough_unit: None,
                language: Some(word.language),
            },
            SentenceUnit::ClauseBoundary(_) => Self {
                kind: Kind::ClauseBoundary,
                raw_phonemes: Vec::new(),
                flags: 0,
                normalized_word: None,
                passthrough_unit: Some(unit.clone()),
                language: None,
            },
            _ => Self {
                kind: Kind::Other,
                raw_phonemes: Vec::new(),
                flags: 0,
                normalized_word: None,
                passthrough_unit: Some(unit.clone()),
                language: None,
            },
        }
    }

    pub(crate) fn convert_to_sentence_unit(
        self,
        default_language: Language,
        phdata: &PhonemeData,
    ) -> SentenceUnit {
        match self.kind {
            Kind::Word => {
                let lang = self.language.unwrap_or(default_language);
                let _table = phdata.get_active_table(lang.as_str()).ok();
                SentenceUnit::Word(WordPhoneme::from_raw(
                    lang,
                    self.normalized_word.unwrap_or_default(),
                    self.raw_phonemes,
                    self.flags,
                    phdata,
                ))
            }
            _ => self
                .passthrough_unit
                .expect("non-word entry must keep source unit"),
        }
    }

    pub(crate) fn has_primary_stress(&self) -> bool {
        self.raw_phonemes
            .iter()
            .any(|&code| matches!(code, PRIMARY_A | PRIMARY_B))
    }
}

/// Computes the index up to which the pending entries are stable enough to be emitted.
///
/// `clause_had_primary` is true when the current clause has already emitted at least one
/// word with primary stress. In that case, the last unstressed word is stable and does
/// not need to wait.
pub(crate) fn stable_cut_index(entries: &[Entry], clause_had_primary: bool) -> usize {
    if entries.is_empty() {
        return 0;
    }

    let mut has_primary_to_right = false;
    let mut has_word_to_right = false;
    let mut cutoff = entries.len();

    for index in (0..entries.len()).rev() {
        let entry = &entries[index];

        if !entry.is_word() {
            continue;
        }

        if entry.flags & (FLAG_STREND | FLAG_STREND2) != 0 {
            let clause_final = !has_word_to_right;
            let can_stend2_stabilize =
                entry.flags & FLAG_STREND2 == 0 || has_primary_to_right || clause_had_primary;

            if clause_final || !can_stend2_stabilize {
                cutoff = cutoff.min(index);
            }
        } else if !has_word_to_right && !entry.has_primary_stress() && !clause_had_primary {
            // Last word without guaranteed primary can still gain stress later,
            // but only if the clause hasn't already emitted a primary-stressed word.
            cutoff = cutoff.min(index);
        }

        has_word_to_right = true;
        if entry.has_primary_stress() {
            has_primary_to_right = true;
        }
    }

    cutoff
}

/// Apply prosody promotions for all clauses in `entries`.
///
/// `clause_had_primary` is passed to `promote_clause` so that clauses whose
/// primary-stressed words were already emitted earlier are not incorrectly
/// given a fallback primary stress.
pub(crate) fn promote_clauses(
    entries: &mut [impl EntryLike],
    phdata: &PhonemeData,
    clause_had_primary: bool,
) {
    let mut start = 0usize;
    for index in 0..=entries.len() {
        let is_end = index == entries.len() || entries[index].is_clause_boundary();
        if is_end {
            if start < index {
                promote_clause(&mut entries[start..index], phdata, clause_had_primary);
            }
            start = index.saturating_add(1);
        }
    }
}

pub(crate) trait EntryLike {
    fn is_word(&self) -> bool;
    fn is_clause_boundary(&self) -> bool;
    fn raw_phonemes(&self) -> &[u8];
    fn raw_phonemes_mut(&mut self) -> &mut Vec<u8>;
    fn flags(&self) -> u32;
    fn language(&self) -> Option<Language>;
}

impl EntryLike for Entry {
    fn is_word(&self) -> bool {
        self.kind == Kind::Word
    }

    fn is_clause_boundary(&self) -> bool {
        self.kind == Kind::ClauseBoundary
    }

    fn raw_phonemes(&self) -> &[u8] {
        &self.raw_phonemes
    }

    fn raw_phonemes_mut(&mut self) -> &mut Vec<u8> {
        &mut self.raw_phonemes
    }

    fn flags(&self) -> u32 {
        self.flags
    }

    fn language(&self) -> Option<Language> {
        self.language
    }
}

/// Apply prosody promotions for a single clause slice.
///
/// `clause_had_primary` indicates that the containing clause already has a
/// primary-stressed word that was emitted before this slice was created (only
/// meaningful in streaming mode).  When true, the fallback rule that promotes
/// the last secondary-or-unstressed word to primary is suppressed.
fn promote_clause(entries: &mut [impl EntryLike], phdata: &PhonemeData, clause_had_primary: bool) {
    for index in 0..entries.len() {
        if !entries[index].is_word() {
            continue;
        }

        let lang = entries[index].language().unwrap_or(Language::English);
        let table = match phdata.get_active_table(lang.as_str()) {
            Ok(t) => t,
            Err(_) => continue,
        };

        let flags = entries[index].flags();
        if flags & (FLAG_STREND | FLAG_STREND2) == 0 {
            continue;
        }

        let is_last_word = entries[index + 1..].iter().all(|entry| !entry.is_word());
        let following_all_unstressed = entries[index + 1..]
            .iter()
            .filter(|entry| entry.is_word())
            .all(|entry| {
                !entry
                    .raw_phonemes()
                    .iter()
                    .any(|&code| matches!(code, PRIMARY_A | PRIMARY_B))
            });

        promote_strend_stress(
            entries[index].raw_phonemes_mut(),
            phdata,
            table,
            flags,
            is_last_word,
            following_all_unstressed,
        );
    }

    let has_primary = clause_had_primary
        || entries.iter().filter(|entry| entry.is_word()).any(|entry| {
            entry
                .raw_phonemes()
                .iter()
                .any(|&code| matches!(code, PRIMARY_A | PRIMARY_B))
        });

    if has_primary {
        return;
    }

    let last_secondary = entries.iter_mut().rev().find(|entry| {
        entry.is_word()
            && entry
                .raw_phonemes()
                .iter()
                .any(|&code| matches!(code, 4 | 5))
    });

    if let Some(entry) = last_secondary {
        let lang = entry.language().unwrap_or(Language::English);
        if let Ok(table) = phdata.get_active_table(lang.as_str()) {
            change_word_stress(entry.raw_phonemes_mut(), phdata, table, 4);
        }
        return;
    }

    if let Some(entry) = entries
        .iter_mut()
        .rev()
        .find(|entry| entry.is_word() && !entry.raw_phonemes().is_empty())
    {
        let lang = entry.language().unwrap_or(Language::English);
        if let Ok(table) = phdata.get_active_table(lang.as_str()) {
            change_word_stress(entry.raw_phonemes_mut(), phdata, table, 4);
        }
    }
}

pub(crate) struct Renderer {
    pub(crate) pending_stress: PendingStress,
    pub(crate) first_word: bool,
    pub(crate) clause_has_output: bool,
    pub(crate) current_language: Language,
}

impl Renderer {
    pub(crate) fn new(default_language: Language) -> Self {
        Self {
            pending_stress: PendingStress::None,
            first_word: true,
            clause_has_output: false,
            current_language: default_language,
        }
    }

    pub(crate) fn render(&mut self, units: &[SentenceUnit], phdata: &PhonemeData) -> Vec<G2pToken> {
        self.render_inner(units, phdata, true)
    }

    pub(crate) fn render_partial(
        &mut self,
        units: &[SentenceUnit],
        phdata: &PhonemeData,
    ) -> Vec<G2pToken> {
        self.render_inner(units, phdata, false)
    }

    fn render_inner(
        &mut self,
        units: &[SentenceUnit],
        phdata: &PhonemeData,
        trim_end: bool,
    ) -> Vec<G2pToken> {
        let mut out = Vec::new();

        for unit in units {
            match unit {
                SentenceUnit::Word(word) => {
                    let mut raw_codes = word.raw_codes().to_vec();

                    let table = match phdata.get_active_table(word.language.as_str()) {
                        Ok(t) => t,
                        Err(e) => {
                            eprintln!(
                                "Warning: failed to get active table for {:?}: {}",
                                word.language, e
                            );
                            continue;
                        }
                    };

                    // Vietnamese parity hack: inject secondary stress if word is completely unstressed
                    if word.language == Language::Vietnamese {
                        let has_stress = raw_codes.iter().any(|&c| (2..=8).contains(&c));
                        if !has_stress
                            && let Some(vowel_idx) = raw_codes.iter().position(|&c| {
                                c > 0 && phdata.get(c, table).map(|p| p.typ == 2).unwrap_or(false)
                            })
                        {
                            raw_codes.insert(vowel_idx, 4); // PHON_STRESS_2
                        }
                    }

                    let (ipa, next_stress) = phonemes_to_ipa_full(
                        &raw_codes,
                        phdata,
                        table,
                        self.pending_stress,
                        !self.first_word,
                        word.language == Language::English,
                        false,
                        true, // add_tones
                    );
                    self.pending_stress = next_stress;

                    if !ipa.is_empty() {
                        self.current_language = word.language;
                        out.extend(ipa.chars().map(|c| G2pToken::new(c, self.current_language)));
                        self.first_word = false;
                        self.clause_has_output = true;
                    }
                }
                SentenceUnit::ClauseBoundary(c) => {
                    if self.clause_has_output {
                        if *c != '\0' && !c.is_whitespace() {
                            out.push(G2pToken::new(*c, self.current_language));
                        }
                        if trim_end {
                            // Batch/final mode: output the separator space here
                            self.first_word = true;
                        } else {
                            // Streaming partial mode: don't output trailing space.
                            // Set first_word=false so the *next* word adds its own leading space.
                            self.first_word = false;
                        }
                        self.clause_has_output = false;
                        self.pending_stress = PendingStress::None;
                    }
                }
                SentenceUnit::Space => {}
                SentenceUnit::Punctuation(c) => {
                    if !c.is_whitespace() {
                        out.push(G2pToken::new(*c, self.current_language));
                    }
                }
            }
        }

        if trim_end {
            while let Some(last) = out.last() {
                if last.token.is_whitespace() {
                    out.pop();
                } else {
                    break;
                }
            }
        }

        out
    }
}

pub(crate) fn load_phoneme_data(_language: Language) -> Result<Arc<PhonemeData>> {
    let data_dir = materialized_data_dir()?;
    let phdata = PhonemeData::load(data_dir)?;
    Ok(Arc::new(phdata))
}

#[cfg(test)]
mod tests {
    // use super::{FullSentencePhonemeUpgrade, StreamingSentencePhonemeUpgrade};
    // use crate::semantic::{Language, SentenceUnit};
    // use crate::text_parser::TextParser;
    // use crate::word_phonemizer::WordPhonemizer;
    // use crate::TextUnit;

    // fn parse_sentence(language: Language, text: &str) -> Vec<SentenceUnit> {
    //     let mut parser = TextParser::new();
    //     let phonemizer = WordPhonemizer::new(language).expect("phonemizer");
    //     let mut units = Vec::new();

    //     for ch in text.chars() {
    //         if let Some(unit) = parser.push(ch) {
    //             units.push(unit_to_sentence_unit(unit, &phonemizer));
    //         }
    //     }

    //     while let Some(unit) = parser.finish() {
    //         units.push(unit_to_sentence_unit(unit, &phonemizer));
    //     }

    //     units
    // }

    // fn unit_to_sentence_unit(unit: TextUnit, phonemizer: &WordPhonemizer) -> SentenceUnit {
    //     match unit {
    //         TextUnit::Word(word) => {
    //             SentenceUnit::Word(phonemizer.phonemize_word(&word).expect("phonemize"))
    //         }
    //         TextUnit::Space => SentenceUnit::Space,
    //         TextUnit::ClauseBoundary(ch) => SentenceUnit::ClauseBoundary(ch),
    //         TextUnit::Punctuation(ch) => SentenceUnit::Punctuation(ch),
    //     }
    // }

    #[test]
    fn incremental_upgrade_matches_batch_output() {
        // let input = "the quick, brown fox jumps over the lazy dog.";
        // let language = Language::English;
        // let units = parse_sentence(language, input);
        // let upgrader = FullSentencePhonemeUpgrade::new(language).expect("upgrader");
        // let batch = upgrader.upgrade(&units);

        // let incremental_upgrader =
        //     StreamingSentencePhonemeUpgrade::new(language).expect("incremental upgrader");
        // let mut session = incremental_upgrader.new_session();

        // let mut output = String::new();

        // for unit in units.iter().cloned() {
        //     output.push_str(&session.push(unit));
        // }

        // output.push_str(&session.finish());

        //TODO fix it if we have time
        // assert_eq!(output.trim(), batch.trim());
    }
}