piper-phoneme-streaming 0.1.1

//! Binary phoneme data loader.
//!
//! Port of `LoadPhData()`, `ReadPhFile()`, and `SelectPhonemeTable()` from
//! `synthdata.c`.  Reads `phontab`, `phonindex`, `phondata`, and
//! `intonations` from the espeak-ng data directory.
//
// Reads four binary files from the espeak-ng-data directory:
//
//   phontab      — phoneme table list (parsed here)
//   phonindex    — u16 offsets into phondata (kept as raw bytes)
//   phondata     — synthesis waveform/formant data (kept as raw bytes)
//   intonations  — array of TUNE structs (kept as raw bytes)

use std::collections::HashMap;
use std::path::Path;

use super::{
    N_PHONEME_TAB, N_PHONEME_TAB_NAME, VERSION_PHDATA,
    table::{PhonemeTab, PhonemeTabList},
};
use crate::error::{Error, Result};

const PHONEME_TAB_ENTRY_SIZE: usize = 16;

/// A pre-calculated mapping of phoneme codes to their positions in the raw tables
/// for a specific language/voice.
#[derive(Clone)]
pub struct ActiveTable {
    /// For each phoneme code (0-255): which (table_idx, entry_idx) is active.
    /// `None` means "no phoneme with that code in this set".
    active: Box<[Option<(usize, usize)>; N_PHONEME_TAB]>,
}

impl ActiveTable {
    fn new() -> Self {
        Self {
            active: Box::new([None; N_PHONEME_TAB]),
        }
    }

    fn fill_from_table(&mut self, phdata: &PhonemeData, idx: usize) {
        let includes = phdata.tables[idx].includes;
        if includes > 0 {
            self.fill_from_table(phdata, (includes - 1) as usize);
        }
        let n = phdata.tables[idx].phonemes.len();
        for entry_idx in 0..n {
            let code = phdata.tables[idx].phonemes[entry_idx].code as usize;
            if code < N_PHONEME_TAB {
                self.active[code] = Some((idx, entry_idx));
            }
        }
    }

    pub fn get_slot(&self, code: u8) -> Option<(usize, usize)> {
        self.active[code as usize]
    }
}

// ---------------------------------------------------------------------------
// PhonemeData — the fully loaded dataset
// ---------------------------------------------------------------------------

/// All phoneme data loaded from an espeak-ng-data directory.
#[derive(Clone)]
pub struct PhonemeData {
    /// All phoneme table definitions parsed from `phontab`.
    pub tables: Vec<PhonemeTabList>,
    /// Sample rate read from `phondata` header (typically 22050).
    pub sample_rate: u32,

    /// Raw bytes of `phondata` (waveform / formant data).
    pub phondata: Vec<u8>,
    /// Raw bytes of `phonindex` (u16 offsets into phondata).
    pub phonindex: Vec<u8>,
    /// Raw bytes of `intonations` (array of TUNE structs, 68 bytes each).
    pub intonations: Vec<u8>,

    /// Pre-calculated active tables for supported languages.
    language_tables: HashMap<String, ActiveTable>,
}

impl PhonemeData {
    // -----------------------------------------------------------------------
    // Loading
    // -----------------------------------------------------------------------

    /// Load all phoneme data from the given data directory.
    pub fn load(data_dir: &Path) -> Result<Self> {
        let phoneme_tab_data = read_file(data_dir, "phontab")?;
        let phonindex = read_file(data_dir, "phonindex")?;
        let phondata = read_file(data_dir, "phondata")?;
        let intonations = read_file(data_dir, "intonations")?;

        if phondata.len() < 8 {
            return Err(Error::InvalidData("phondata too short".into()));
        }
        let version = u32::from_le_bytes(phondata[0..4].try_into().unwrap());
        if version != VERSION_PHDATA {
            return Err(Error::VersionMismatch {
                got: version,
                expected: VERSION_PHDATA,
            });
        }
        let sample_rate = u32::from_le_bytes(phondata[4..8].try_into().unwrap());

        let tables = parse_phontab(&phoneme_tab_data)?;

        let mut me = Self {
            tables,
            sample_rate,
            phondata,
            phonindex,
            intonations,
            language_tables: HashMap::new(),
        };

        // Pre-calculate tables for supported languages (en, vi)
        for lang in &["en", "vi"] {
            if let Ok(idx) = me.find_table(lang) {
                let mut at = ActiveTable::new();
                at.fill_from_table(&me, idx);
                me.language_tables.insert(lang.to_string(), at);
            }
        }

        Ok(me)
    }

    /// Get the pre-calculated active table for a language.
    pub fn get_active_table(&self, lang_name: &str) -> Result<&ActiveTable> {
        self.language_tables.get(lang_name).ok_or_else(|| {
            Error::InvalidData(format!(
                "Phoneme table for language '{}' not pre-calculated",
                lang_name
            ))
        })
    }

    pub fn find_table(&self, name: &str) -> Result<usize> {
        self.tables
            .iter()
            .position(|t| t.name == name)
            .ok_or_else(|| Error::InvalidData(format!("phoneme table '{name}' not found")))
    }

    /// Backwards compatibility: select a table by name.
    /// WARNING: This is now a slow operation as it re-calculates the active table.
    /// Prefer using get_active_table().
    pub fn select_table_by_name(&mut self, name: &str) -> Result<usize> {
        let idx = self.find_table(name)?;
        let mut at = ActiveTable::new();
        at.fill_from_table(self, idx);
        // We don't store this 'current' table anymore, but we return the index
        Ok(idx)
    }

    // -----------------------------------------------------------------------
    // Phoneme lookup
    // -----------------------------------------------------------------------

    pub fn phoneme_code(&self, mnem: u32, table: &ActiveTable) -> u8 {
        for slot in table.active.iter().flatten() {
            let (table_idx, entry_idx) = *slot;
            let ph = &self.tables[table_idx].phonemes[entry_idx];
            if ph.mnemonic == mnem {
                return ph.code;
            }
        }
        0
    }

    pub fn lookup_phoneme(&self, name: &str, table: &ActiveTable) -> u8 {
        self.phoneme_code(PhonemeTab::pack_mnemonic(name), table)
    }

    pub fn get(&self, code: u8, table: &ActiveTable) -> Option<&PhonemeTab> {
        let (table_idx, entry_idx) = table.get_slot(code)?;
        Some(&self.tables[table_idx].phonemes[entry_idx])
    }

    /// Resolve a phoneme code through synthesis-stage `ChangeIf` instructions.
    pub fn resolve_stressed_phoneme(&self, code: u8, is_stressed: bool, table: &ActiveTable) -> u8 {
        let Some(ph) = self.get(code, table) else {
            return code;
        };
        if ph.program == 0 {
            return code;
        }

        // STRESS constants (from synthesize.h)
        const STRESS_IS_DIMINISHED: u8 = 0;
        const STRESS_IS_PRIMARY: u8 = 4;

        // condition_level[condition] table from StressCondition()
        // condition 0→1, 1→2, 2→4(PRIMARY), 3→15
        const CONDITION_LEVEL: [u8; 4] = [1, 2, 4, 15];

        let stress_level: u8 = if is_stressed {
            STRESS_IS_PRIMARY
        } else {
            STRESS_IS_DIMINISHED
        };
        let prog = ph.program as usize;
        let pi = &self.phonindex;

        if ph.typ != 2 {
            return code;
        } // not a vowel

        const THIS_PH_IS_MAX_STRESS: u16 = 0x2884; // thisPh(isMaxStress)

        let mut i = 0usize;
        while i < 16 {
            let off = (prog + i) * 2;
            if off + 2 > pi.len() {
                break;
            }
            let w = u16::from_le_bytes([pi[off], pi[off + 1]]);
            let instn_type = w >> 12;
            let instn2 = ((w >> 8) & 0xf) as u8;
            let data_u8 = (w & 0xff) as u8; // lower 8 bits for phoneme code / jump offset

            if instn_type == 1 && instn2 < 8 {
                let fires = if instn2 == STRESS_IS_PRIMARY {
                    is_stressed
                } else if (instn2 as usize) < CONDITION_LEVEL.len() {
                    stress_level < CONDITION_LEVEL[instn2 as usize]
                } else {
                    false
                };

                if fires && data_u8 != 0 {
                    return data_u8; // changed phoneme code
                }
                i += 1;
            } else if w == THIS_PH_IS_MAX_STRESS {
                if !is_stressed {
                    let next_off = (prog + i + 1) * 2;
                    if next_off + 2 <= pi.len() {
                        let jw = u16::from_le_bytes([pi[next_off], pi[next_off + 1]]);
                        if (jw & 0xf800) == 0x6800 {
                            let jump = (jw & 0xff) as usize;
                            i += 2 + jump;
                            continue;
                        }
                    }
                    break;
                } else {
                    let next_off = (prog + i + 1) * 2;
                    if next_off + 2 <= pi.len() {
                        let jw = u16::from_le_bytes([pi[next_off], pi[next_off + 1]]);
                        if (jw & 0xf800) == 0x6800 {
                            i += 2;
                            continue;
                        }
                    }
                    i += 1;
                }
            } else if instn_type == 2 || instn_type == 3 {
                let next_off = (prog + i + 1) * 2;
                if next_off + 2 <= pi.len() {
                    let jw = u16::from_le_bytes([pi[next_off], pi[next_off + 1]]);
                    if (jw & 0xf800) == 0x6800 {
                        let jump = (jw & 0xff) as usize;
                        if !is_stressed {
                            i += 2 + jump;
                        } else {
                            i += 2;
                        }
                        continue;
                    }
                }
                break;
            } else if instn_type == 6 {
                if (instn2 >> 1) == 0 {
                    let jump_by = (data_u8 as usize).saturating_sub(1);
                    i += 1 + jump_by;
                } else {
                    break;
                }
            } else if instn_type == 0 {
                if instn2 == 1 {
                    return data_u8;
                }
                i += 1;
            } else if instn_type >= 0xb {
                break;
            } else {
                i += 1;
            }
        }
        code
    }

    pub fn phoneme_ipa_string(&self, program: u32) -> Option<String> {
        if program == 0 {
            return None;
        }

        const I_IPA_NAME: u16 = 0x0d;
        const MAX_SCAN: usize = 8;

        let phonindex = &self.phonindex;
        let prog = program as usize;

        let first_offset = prog * 2;
        if first_offset + 2 > phonindex.len() {
            return None;
        }
        let first_instn =
            u16::from_le_bytes([phonindex[first_offset], phonindex[first_offset + 1]]);

        let max_scan = if first_instn >= 0xb000 { 1 } else { MAX_SCAN };

        for i in 0..max_scan {
            let offset = (prog + i) * 2;
            if offset + 2 > phonindex.len() {
                break;
            }

            let instn = u16::from_le_bytes([phonindex[offset], phonindex[offset + 1]]);
            let instn_type = instn >> 12;
            let instn2 = (instn >> 8) & 0xf;
            let data = (instn & 0xff) as usize;

            if instn_type == 2 || instn_type == 3 || instn_type == 6 {
                return None;
            }
            if instn == 0x9100 {
                return None;
            }

            if i > 0 && instn >= 0xb000 {
                return None;
            }

            if instn_type == 0 && instn2 as u16 == I_IPA_NAME {
                if data == 0 {
                    return None;
                }
                let mut ipa_bytes = Vec::with_capacity(data);
                let n_words = (data + 1) / 2;
                for j in 0..n_words {
                    let word_off = (prog + i + 1 + j) * 2;
                    if word_off + 2 > phonindex.len() {
                        break;
                    }
                    let word = u16::from_le_bytes([phonindex[word_off], phonindex[word_off + 1]]);
                    ipa_bytes.push(((word >> 8) & 0xff) as u8);
                    ipa_bytes.push((word & 0xff) as u8);
                }
                ipa_bytes.truncate(data);
                return String::from_utf8(ipa_bytes)
                    .ok()
                    .filter(|s| !s.is_empty() && !s.starts_with('\u{0001}'));
            }
        }
        None
    }

    // -----------------------------------------------------------------------
    // Helpers
    // -----------------------------------------------------------------------

    pub fn n_tables(&self) -> usize {
        self.tables.len()
    }

    pub fn n_tunes(&self) -> usize {
        self.intonations.len() / 68
    }

    pub fn phondata_at(&self, offset: usize) -> &[u8] {
        &self.phondata[offset..]
    }
}

// ---------------------------------------------------------------------------
// phontab parser
// ---------------------------------------------------------------------------

fn parse_phontab(data: &[u8]) -> Result<Vec<PhonemeTabList>> {
    if data.is_empty() {
        return Err(Error::InvalidData("phontab is empty".into()));
    }

    let n_tables = data[0] as usize;
    let mut pos = 4usize; // skip [n_tables, 0, 0, 0]
    let mut tables = Vec::with_capacity(n_tables);

    for _i in 0..n_tables {
        if pos + 4 > data.len() {
            return Err(Error::InvalidData(
                "phontab truncated in table header".into(),
            ));
        }
        let n_phonemes = data[pos] as usize;
        let includes = data[pos + 1];
        pos += 4;

        if pos + N_PHONEME_TAB_NAME > data.len() {
            return Err(Error::InvalidData("phontab truncated in table name".into()));
        }
        let name_buf: &[u8; N_PHONEME_TAB_NAME] =
            data[pos..pos + N_PHONEME_TAB_NAME].try_into().unwrap();
        let name = PhonemeTabList::parse_name(name_buf);
        pos += N_PHONEME_TAB_NAME;

        let entries_size = n_phonemes * PHONEME_TAB_ENTRY_SIZE;
        if pos + entries_size > data.len() {
            return Err(Error::InvalidData(format!(
                "phontab truncated in phoneme entries for table '{name}'"
            )));
        }
        let mut phonemes = Vec::with_capacity(n_phonemes);
        for j in 0..n_phonemes {
            let off = pos + j * PHONEME_TAB_ENTRY_SIZE;
            let entry: &[u8; 16] = data[off..off + 16].try_into().unwrap();
            phonemes.push(PhonemeTab::from_bytes(entry));
        }
        pos += entries_size;

        tables.push(PhonemeTabList {
            name,
            phonemes,
            n_phonemes,
            includes,
        });
    }

    Ok(tables)
}

fn read_file(dir: &Path, name: &str) -> Result<Vec<u8>> {
    let path = dir.join(name);
    std::fs::read(&path).map_err(|e| Error::Io(e))
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    const DATA_DIR: &str = "/usr/share/espeak-ng-data";

    fn data_available() -> bool {
        Path::new(DATA_DIR).join("phontab").exists()
    }

    #[test]
    fn load_phdata_basic() {
        if !data_available() {
            return;
        }
        let d = PhonemeData::load(Path::new(DATA_DIR)).expect("load_phdata");
        assert!(d.n_tables() >= 130);
        assert_eq!(d.sample_rate, 22050);
        assert!(d.n_tunes() >= 30);
        assert!(d.phondata.len() > 8);
    }

    #[test]
    fn table_names_include_base() {
        if !data_available() {
            return;
        }
        let d = PhonemeData::load(Path::new(DATA_DIR)).unwrap();
        assert_eq!(d.tables[0].name, "base");
        assert_eq!(d.tables[1].name, "base1");
        assert!(
            d.tables.iter().any(|t| t.name == "en"),
            "no 'en' table found"
        );
    }

    #[test]
    fn base_table_phoneme_count() {
        if !data_available() {
            return;
        }
        let d = PhonemeData::load(Path::new(DATA_DIR)).unwrap();
        assert_eq!(d.tables[0].n_phonemes, 35, "base table phoneme count");
        assert_eq!(d.tables[0].includes, 0, "base table has no parent");
    }

    #[test]
    fn base1_includes_base() {
        if !data_available() {
            return;
        }
        let d = PhonemeData::load(Path::new(DATA_DIR)).unwrap();
        assert_eq!(d.tables[1].includes, 1);
    }

    #[test]
    fn get_active_table_en() {
        if !data_available() {
            return;
        }
        let d = PhonemeData::load(Path::new(DATA_DIR)).unwrap();
        let table = d
            .get_active_table("en")
            .expect("'en' table should be pre-calculated");
        let code = d.lookup_phoneme("t", table);
        assert!(code > 0);
    }

    #[test]
    fn lookup_pause_phoneme() {
        if !data_available() {
            return;
        }
        let d = PhonemeData::load(Path::new(DATA_DIR)).unwrap();
        let table = d.get_active_table("en").unwrap();
        let code = d.lookup_phoneme("_", table);
        assert_eq!(code, 10, "pause phoneme code");
    }

    #[test]
    fn lookup_unknown_returns_zero() {
        if !data_available() {
            return;
        }
        let d = PhonemeData::load(Path::new(DATA_DIR)).unwrap();
        let table = d.get_active_table("en").unwrap();
        assert_eq!(d.lookup_phoneme("???", table), 0);
    }

    #[test]
    fn get_nonexistent_table_errors() {
        if !data_available() {
            return;
        }
        let d = PhonemeData::load(Path::new(DATA_DIR)).unwrap();
        assert!(d.get_active_table("no_such_language").is_err());
    }

    #[test]
    fn find_table_returns_index() {
        if !data_available() {
            return;
        }
        let d = PhonemeData::load(Path::new(DATA_DIR)).unwrap();
        let idx = d.find_table("base").unwrap();
        assert_eq!(idx, 0);
    }
}