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();
    }

    while result.contains("  ") {
        result = result.replace("  ", " ");
    }

    result.trim().to_string()
}

/// Replacement pairs for American → British English conversion.
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 (word-boundary needed)
    ("\\bcolor\\b", "colour"),
    ("\\bfavor\\b", "favour"),
    ("\\bhonor\\b", "honour"),
    ("\\bhumor\\b", "humour"),
    ("\\blabor\\b", "labour"),
    ("\\bneighbor\\b", "neighbour"),
    ("behavior", "behaviour"),
    // -er → -re (word-boundary to avoid "parameter" etc.)
    ("\\bcenter\\b", "centre"),
    ("\\bfiber\\b", "fibre"),
    ("\\bliter\\b", "litre"),
    ("\\bmeter\\b", "metre"),
    ("\\btheater\\b", "theatre"),
    // -ense → -ence
    ("defense", "defence"),
    ("offense", "offence"),
    // Other
    ("\\bcatalog\\b", "catalogue"),
    ("\\bdialog\\b", "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 {
        let pattern = if us.contains("\\b") {
            format!("(?i){us}")
        } else {
            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 us_base = us.replace("\\b", "");
                    let base_len = us_base.len();
                    let suffix = if matched.len() > base_len {
                        &matched[base_len..]
                    } else {
                        ""
                    };
                    let Some(first_char) = matched.chars().next() else {
                        return String::new();
                    };
                    let uk_clean = uk.replace("\\b", "");
                    if first_char.is_uppercase() {
                        let mut chars = uk_clean.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_clean}{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."));
    }
}