Lumotia/crates/ai-formatting/src/rule_based.rs

use std::sync::LazyLock;

/// Compiled filler word regexes (built once, reused across calls).
/// Uses \b word boundaries instead of lookbehinds (regex-lite does not
/// support lookaround assertions).
static FILLER_REGEXES: LazyLock<Vec<regex_lite::Regex>> = LazyLock::new(|| {
    let fillers = [
        "um",
        "uh",
        "er",
        "ah",
        "like",
        "you know",
        "sort of",
        "kind of",
        "I mean",
        "basically",
        "actually",
        "literally",
    ];
    fillers
        .iter()
        .filter_map(|filler| {
            let escaped = regex_lite::escape(filler);
            let pattern = format!(r"(?i)\b{escaped}\b[,.]?\s*");
            regex_lite::Regex::new(&pattern).ok()
        })
        .collect()
});

/// Remove common filler words from transcription text (case-insensitive).
pub fn remove_fillers(text: &str) -> String {
    let mut result = text.to_string();

    for re in FILLER_REGEXES.iter() {
        result = re.replace_all(&result, " ").to_string();
    }

    // Collapse runs of whitespace in a single pass.
    let mut collapsed = String::with_capacity(result.len());
    let mut prev_space = false;
    for ch in result.chars() {
        if ch == ' ' {
            if !prev_space {
                collapsed.push(' ');
            }
            prev_space = true;
        } else {
            prev_space = false;
            collapsed.push(ch);
        }
    }

    collapsed.trim().to_string()
}

/// Replacement pairs for American to British English conversion.
///
/// All entries are plain base words (no regex metacharacters). The
/// `to_british_english` function wraps every entry with `\b` word
/// boundaries and optional suffix matching automatically.
static BRITISH_REPLACEMENTS: &[(&str, &str)] = &[
    // -ize → -ise (and inflected forms)
    ("organize", "organise"),
    ("recognize", "recognise"),
    ("realize", "realise"),
    ("analyze", "analyse"),
    ("apologize", "apologise"),
    ("authorize", "authorise"),
    ("categorize", "categorise"),
    ("characterize", "characterise"),
    ("customize", "customise"),
    ("digitize", "digitise"),
    ("emphasize", "emphasise"),
    ("finalize", "finalise"),
    ("generalize", "generalise"),
    ("harmonize", "harmonise"),
    ("initialize", "initialise"),
    ("maximize", "maximise"),
    ("minimize", "minimise"),
    ("modernize", "modernise"),
    ("normalize", "normalise"),
    ("optimize", "optimise"),
    ("prioritize", "prioritise"),
    ("revolutionize", "revolutionise"),
    ("specialize", "specialise"),
    ("standardize", "standardise"),
    ("summarize", "summarise"),
    ("utilize", "utilise"),
    // -or → -our
    ("color", "colour"),
    ("favor", "favour"),
    ("honor", "honour"),
    ("humor", "humour"),
    ("labor", "labour"),
    ("neighbor", "neighbour"),
    ("behavior", "behaviour"),
    // -er → -re
    ("center", "centre"),
    ("fiber", "fibre"),
    ("liter", "litre"),
    ("meter", "metre"),
    ("theater", "theatre"),
    // -ense → -ence
    ("defense", "defence"),
    ("offense", "offence"),
    // Other
    ("catalog", "catalogue"),
    ("dialog", "dialogue"),
];

/// Convert American English spelling to British English (word-boundary aware).
pub fn to_british_english(text: &str) -> String {
    let mut result = text.to_string();

    for (us, uk) in BRITISH_REPLACEMENTS {
        // Every entry in BRITISH_REPLACEMENTS is a plain ASCII base word.
        // We wrap it with \b boundaries and optional suffix matching here.
        let pattern = format!("(?i)\\b{}(?:d|s|r|rs)?\\b", regex_lite::escape(us));

        if let Ok(re) = regex_lite::Regex::new(&pattern) {
            result = re
                .replace_all(&result, |caps: &regex_lite::Captures| {
                    let Some(m) = caps.get(0) else {
                        return String::new();
                    };
                    let matched = m.as_str();
                    let base_len = us.len();
                    // SAFETY: byte indexing is correct here because both
                    // the US base word and the suffix characters (d, s, r)
                    // are guaranteed ASCII by the BRITISH_REPLACEMENTS table.
                    debug_assert!(us.is_ascii(), "BRITISH_REPLACEMENTS entries must be ASCII");
                    debug_assert!(matched.is_ascii(), "matched text expected to be ASCII");
                    let suffix = if matched.len() > base_len {
                        &matched[base_len..]
                    } else {
                        ""
                    };
                    let Some(first_char) = matched.chars().next() else {
                        return String::new();
                    };
                    if first_char.is_uppercase() {
                        let mut chars = uk.chars();
                        let Some(first_uk) = chars.next() else {
                            return String::new();
                        };
                        let upper_first: String = first_uk.to_uppercase().collect();
                        format!("{}{}{}", upper_first, chars.collect::<String>(), suffix)
                    } else {
                        format!("{uk}{suffix}")
                    }
                })
                .to_string();
        }
    }

    result
}

/// Basic formatting: capitalise sentences, fix spacing, clean punctuation.
pub fn format_text(text: &str) -> String {
    if text.is_empty() {
        return String::new();
    }

    let mut result = String::with_capacity(text.len());
    let mut capitalise_next = true;

    let chars: Vec<char> = text.chars().collect();
    let mut i = 0;

    while i < chars.len() {
        let c = chars[i];

        if c == ' ' && i + 1 < chars.len() && chars[i + 1] == ' ' {
            i += 1;
            continue;
        }

        if capitalise_next && c.is_alphabetic() {
            result.extend(c.to_uppercase());
            capitalise_next = false;
        } else {
            result.push(c);
        }

        if c == '.' || c == '!' || c == '?' {
            capitalise_next = true;
        }
        if c == '\n' {
            capitalise_next = true;
        }

        i += 1;
    }

    result
}

/// Known hallucination markers that should be filtered from transcriptions.
static HALLUCINATION_MARKERS: &[&str] = &["[blank_audio]", "[music]", "[silence]"];

static AUTO_THANKS_PHRASES: &[&str] = &["thank you.", "thanks.", "you.", "thank you for watching."];

/// Returns true if a segment's text looks like a hallucination.
pub fn is_hallucination(text: &str) -> bool {
    let trimmed = text.trim().to_lowercase();
    if trimmed.is_empty() {
        return true;
    }
    for marker in HALLUCINATION_MARKERS {
        if trimmed.contains(marker) {
            return true;
        }
    }
    if trimmed.len() < 15 {
        for phrase in AUTO_THANKS_PHRASES {
            if trimmed == *phrase {
                return true;
            }
        }
    }
    false
}

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

    #[test]
    fn remove_fillers_strips_um_and_uh() {
        let input = "So um I was thinking uh about this";
        let result = remove_fillers(input);
        assert!(!result.contains("um"));
        assert!(!result.contains("uh"));
    }

    #[test]
    fn remove_fillers_preserves_legitimate_words() {
        let input = "The umbrella was actually useful";
        let result = remove_fillers(input);
        assert!(result.contains("umbrella"));
        assert!(result.contains("useful"));
    }

    #[test]
    fn to_british_english_converts_ize_to_ise() {
        assert!(to_british_english("organize").contains("organise"));
        assert!(to_british_english("realize").contains("realise"));
    }

    #[test]
    fn to_british_english_preserves_case() {
        let result = to_british_english("Organize the files");
        assert!(result.starts_with("Organise"));
    }

    #[test]
    fn to_british_english_handles_colour() {
        assert!(to_british_english("the color is red").contains("colour"));
    }

    #[test]
    fn format_text_capitalises_after_full_stops() {
        let result = format_text("hello world. this is a test");
        assert!(result.starts_with('H'));
        assert!(result.contains(". T"));
    }

    #[test]
    fn format_text_handles_empty_string() {
        assert_eq!(format_text(""), "");
    }

    #[test]
    fn is_hallucination_detects_blank_audio() {
        assert!(is_hallucination("[blank_audio]"));
        assert!(is_hallucination("  [music]  "));
    }

    #[test]
    fn is_hallucination_detects_auto_thanks() {
        assert!(is_hallucination("Thank you."));
        assert!(is_hallucination("thanks."));
    }

    #[test]
    fn is_hallucination_allows_real_text() {
        assert!(!is_hallucination("The meeting is at three o'clock."));
    }
}