tokmat 0.3.3

//! Tokenization helpers for the wanParser Rust port.

use crate::token_model::TokenModel;
use crate::word_definition::tokenizer_boundary;
use pcre2::bytes::{Regex as Pcre2Regex, RegexBuilder as Pcre2RegexBuilder};
use std::collections::{HashMap, HashSet};
use std::hash::BuildHasher;

pub use crate::token_model::{
    TokenClassList, TokenDefinition, load_token_class_list, load_token_definitions,
};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TokenizedResult {
    pub raw_value: String,
    pub tokens: Vec<String>,
    pub types: Vec<String>,
    pub classes: Vec<String>,
}

/// Split an input string using the wanParser word-boundary definition.
///
/// # Panics
///
/// Panics if the built-in PCRE2 word-boundary regex cannot execute, which would indicate an
/// internal invariant violation because the pattern is compiled during crate initialization.
#[must_use]
pub fn split_input_tokens(input: &str) -> Vec<String> {
    split_input_tokens_with(input, &tokenizer_boundary())
}

/// Split using an explicit word-boundary regex instead of the process-global one.
///
/// This is the per-model entry point: a [`TokenModel`] holds its own compiled
/// boundary (from its `WORDDEFINITION.param`), so tokenization is deterministic
/// per model and does not depend on (or race) the process-global word definition.
///
/// # Panics
///
/// Panics if the supplied word-boundary regex cannot execute.
#[must_use]
pub fn split_input_tokens_with(input: &str, boundary_re: &Pcre2Regex) -> Vec<String> {
    let mut tokens = Vec::new();
    let mut segment_start = 0_usize;

    for boundary in boundary_re.find_iter(input.as_bytes()) {
        let boundary = boundary.expect("word boundary regex should execute");
        let boundary_index = boundary.start();
        if boundary_index > segment_start {
            tokens.push(input[segment_start..boundary_index].to_string());
        }
        segment_start = boundary_index;
    }

    if segment_start < input.len() {
        tokens.push(input[segment_start..].to_string());
    }

    tokens
}

/// ASCII whitespace per Unicode `White_Space` (matches `char::is_whitespace`
/// for ASCII, including vertical tab `0x0B` unlike `u8::is_ascii_whitespace`).
const fn ascii_is_whitespace(byte: u8) -> bool {
    matches!(byte, b' ' | b'\t' | b'\n' | 0x0B | 0x0C | b'\r')
}

/// Fast-path classifier mirroring the Python tokenizer shortcuts.
///
/// Computes every character-class predicate in a single pass over the token's
/// `chars` rather than re-walking it ~10 times. All fast-path classes require an
/// ASCII token: this matches the canonical Python wanParser reference, where an
/// accented token (e.g. `Étg`) keeps its raw form as the type rather than being
/// folded into `ALPHA`. French-Canadian handling is provided by the token-class
/// dictionary and upstream accent normalization, not by this classifier.
#[must_use]
pub fn get_token_fast_classifier<S: BuildHasher>(
    token: &str,
    available_names: &HashSet<String, S>,
) -> Option<String> {
    if token.is_empty() {
        return None;
    }

    let mut all_digit = true;
    let mut all_alpha = true;
    let mut has_digit = false;
    let mut has_alpha = false;
    let mut all_digit_or_dash = true;
    let mut all_alpha_dash_apos = true;
    let mut all_alnum = true;
    let mut all_alnum_dash_apos = true;
    let mut is_ascii = true;

    // Compacted (whitespace/'-'/'_' stripped) ASCII characters for the A1A1A1
    // POSTALCODE shape; capped at 6 so a 7th char cheaply rejects.
    let track_postal = available_names.contains("POSTALCODE");
    let mut postal = [0_u8; 6];
    let mut postal_len = 0_usize;
    let mut postal_overflow = false;

    for character in token.chars() {
        let alpha = character.is_alphabetic();
        let digit = character.is_ascii_digit();
        let dash = character == '-';
        let apos = character == '\'';
        let alnum = alpha | digit;
        let ascii = character.is_ascii();

        has_alpha |= alpha;
        has_digit |= digit;
        all_digit &= digit;
        all_alpha &= alpha;
        all_digit_or_dash &= digit | dash;
        all_alpha_dash_apos &= alpha | dash | apos;
        all_alnum &= alnum;
        all_alnum_dash_apos &= alnum | dash | apos;
        is_ascii &= ascii;

        if track_postal && ascii && !postal_overflow {
            let byte = character as u8;
            if byte != b'-' && byte != b'_' && !ascii_is_whitespace(byte) {
                if postal_len < postal.len() {
                    postal[postal_len] = byte;
                    postal_len += 1;
                } else {
                    postal_overflow = true;
                }
            }
        }
    }

    if track_postal
        && is_ascii
        && !postal_overflow
        && postal_len == 6
        && postal[0].is_ascii_alphabetic()
        && postal[1].is_ascii_digit()
        && postal[2].is_ascii_alphabetic()
        && postal[3].is_ascii_digit()
        && postal[4].is_ascii_alphabetic()
        && postal[5].is_ascii_digit()
    {
        return Some("POSTALCODE".to_string());
    }
    if all_digit && available_names.contains("NUM") {
        return Some("NUM".to_string());
    }
    if is_ascii && all_alpha && available_names.contains("ALPHA") {
        return Some("ALPHA".to_string());
    }
    if all_digit_or_dash && has_digit && available_names.contains("NUM_EXTENDED") {
        return Some("NUM_EXTENDED".to_string());
    }
    if is_ascii && all_alpha_dash_apos && has_alpha && available_names.contains("ALPHA_EXTENDED") {
        return Some("ALPHA_EXTENDED".to_string());
    }
    if is_ascii && all_alnum_dash_apos && has_alpha && has_digit {
        if all_alnum && available_names.contains("ALPHA_NUM") {
            return Some("ALPHA_NUM".to_string());
        }
        if available_names.contains("ALPHA_NUM_EXTENDED") {
            return Some("ALPHA_NUM_EXTENDED".to_string());
        }
    }

    None
}

/// Tokenize and classify a cleaned wanParser string.
///
/// # Panics
///
/// Panics if a token definition contains a regex pattern that compiled successfully when the
/// model was loaded but cannot be recompiled with start/end anchors applied here.
#[must_use]
pub fn tokenize_and_classify(
    raw_value: &str,
    token_definitions: &TokenDefinition,
    token_class_list: Option<&TokenClassList>,
) -> TokenizedResult {
    let tokens = split_input_tokens(raw_value);
    let available_names: HashSet<String> = token_definitions
        .iter()
        .map(|(name, _)| name.clone())
        .collect();
    let compiled_patterns: Vec<(String, Pcre2Regex)> = token_definitions
        .iter()
        .map(|(name, pattern)| {
            let anchored = if pattern.starts_with('^') && pattern.ends_with('$') {
                pattern.clone()
            } else {
                format!(
                    "^{}$",
                    pattern.trim_start_matches('^').trim_end_matches('$')
                )
            };
            (name.clone(), compile_token_regex(&anchored))
        })
        .collect();

    let token_class_lookup = build_token_class_lookup(token_class_list);
    let mut types = Vec::with_capacity(tokens.len());
    let mut classes = Vec::with_capacity(tokens.len());

    for token in &tokens {
        let token_type = get_token_fast_classifier(token, &available_names).unwrap_or_else(|| {
            compiled_patterns
                .iter()
                .find_map(|(name, regex)| {
                    regex
                        .is_match(token.as_bytes())
                        .ok()
                        .and_then(|matched| matched.then(|| name.clone()))
                })
                .unwrap_or_else(|| token.clone())
        });

        types.push(token_type.clone());

        if token_class_list.is_some() {
            if token.chars().all(char::is_whitespace) {
                classes.push(token.clone());
            } else {
                classes.push(token_class_lookup.get(token).cloned().unwrap_or(token_type));
            }
        }
    }

    TokenizedResult {
        raw_value: raw_value.to_string(),
        tokens,
        types,
        classes,
    }
}

/// Tokenize using a precompiled [`TokenModel`].
#[must_use]
pub fn tokenize_with_model(raw_value: &str, model: &TokenModel) -> TokenizedResult {
    let tokens = split_input_tokens_with(raw_value, model.word_boundary());
    let mut types = Vec::with_capacity(tokens.len());
    let mut classes = Vec::with_capacity(tokens.len());

    for token in &tokens {
        let token_type =
            get_token_fast_classifier(token, model.available_names()).unwrap_or_else(|| {
                model
                    .compiled_patterns()
                    .iter()
                    .find_map(|(name, regex)| {
                        regex
                            .is_match(token.as_bytes())
                            .ok()
                            .and_then(|matched| matched.then(|| name.clone()))
                    })
                    .unwrap_or_else(|| token.clone())
            });

        types.push(token_type.clone());
        if token.chars().all(char::is_whitespace) {
            classes.push(token.clone());
        } else {
            classes.push(
                model
                    .token_class_lookup()
                    .get(token)
                    .cloned()
                    .unwrap_or(token_type),
            );
        }
    }

    TokenizedResult {
        raw_value: raw_value.to_string(),
        tokens,
        types,
        classes,
    }
}

fn build_token_class_lookup(token_class_list: Option<&TokenClassList>) -> HashMap<String, String> {
    let mut temp_lookup: HashMap<String, Vec<String>> = HashMap::new();
    if let Some(class_list) = token_class_list {
        for (class_name, values) in class_list {
            for value in values {
                temp_lookup
                    .entry(value.clone())
                    .or_default()
                    .push(class_name.clone());
            }
        }
    }

    temp_lookup
        .into_iter()
        .map(|(value, classes)| (value, classes.join("|")))
        .collect()
}

fn compile_token_regex(pattern: &str) -> Pcre2Regex {
    Pcre2RegexBuilder::new()
        .utf(true)
        .ucp(true)
        .jit_if_available(true)
        .build(pattern)
        .expect("valid token regex")
}

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

    #[test]
    fn test_split_input_tokens_preserves_extended_boundaries() {
        // Shared: depends on the default global word boundary; exclude the
        // definition-mutating test from running concurrently.
        let _guard = crate::word_definition::WORD_DEF_TEST_LOCK
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        assert_eq!(
            split_input_tokens("123 MAIN ST"),
            vec!["123", " ", "MAIN", " ", "ST"]
        );
        assert_eq!(
            split_input_tokens("APT-210 O'CONNOR"),
            vec!["APT-210", " ", "O'CONNOR"]
        );
        assert_eq!(
            split_input_tokens("WORD--ANOTHER...END"),
            vec!["WORD--ANOTHER", "...", "END"]
        );
    }

    #[test]
    fn test_get_token_fast_classifier_handles_common_shapes() {
        let names: HashSet<_> = vec![
            "NUM",
            "ALPHA",
            "POSTALCODE",
            "ALPHA_EXTENDED",
            "ALPHA_NUM_EXTENDED",
        ]
        .into_iter()
        .map(String::from)
        .collect();
        assert_eq!(
            get_token_fast_classifier("123", &names),
            Some("NUM".to_string())
        );
        assert_eq!(
            get_token_fast_classifier("MAIN", &names),
            Some("ALPHA".to_string())
        );
        assert_eq!(
            get_token_fast_classifier("K1A0B1", &names),
            Some("POSTALCODE".to_string())
        );
        assert_eq!(
            get_token_fast_classifier("O'CONNOR", &names),
            Some("ALPHA_EXTENDED".to_string())
        );
        assert_eq!(
            get_token_fast_classifier("APT-210", &names),
            Some("ALPHA_NUM_EXTENDED".to_string())
        );
    }

    #[test]
    fn test_tokenize_and_classify_absorbs_wordlike_hyphen_into_extended_classes() {
        // The default word definition includes hyphen, so the tokenizer never
        // emits a standalone "-" token for these shapes. TEL therefore sees the
        // extended class as one segment, not a class stream containing "-".
        let _guard = crate::word_definition::WORD_DEF_TEST_LOCK
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let defs = vec![
            ("NUM".to_string(), r"^\d+$".to_string()),
            ("ALPHA".to_string(), r"^[A-Z]+$".to_string()),
            ("NUM_EXTENDED".to_string(), r"^[\d-]+$".to_string()),
            ("ALPHA_EXTENDED".to_string(), r"^[A-Z'-]+$".to_string()),
            (
                "ALPHA_NUM_EXTENDED".to_string(),
                r"(?=.*[0-9])(?=.*[A-Z])[A-Z0-9'-]+$".to_string(),
            ),
        ];
        let classes = Vec::new();
        let tokenized = tokenize_and_classify("11-47 OAK-VIEW A-12", &defs, Some(&classes));

        assert_eq!(
            tokenized.tokens,
            vec!["11-47", " ", "OAK-VIEW", " ", "A-12"]
        );
        assert_eq!(
            tokenized.types,
            vec![
                "NUM_EXTENDED",
                " ",
                "ALPHA_EXTENDED",
                " ",
                "ALPHA_NUM_EXTENDED"
            ]
        );
        assert_eq!(
            tokenized.classes,
            vec![
                "NUM_EXTENDED",
                " ",
                "ALPHA_EXTENDED",
                " ",
                "ALPHA_NUM_EXTENDED"
            ]
        );
        assert!(
            !tokenized.classes.iter().any(|class| class == "-"),
            "hyphen is inside the extended token and must not appear as its own class"
        );
    }

    #[test]
    fn test_get_token_fast_classifier_accented_tokens_fall_through() {
        // Canonical wanParser parity: accented (non-ASCII) tokens are NOT
        // folded into ALPHA/ALPHA_EXTENDED by the fast path -- they fall through
        // (return None) so they keep their raw form as the type. French-Canadian
        // handling comes from the class dictionary / upstream normalization.
        let names: HashSet<_> = vec!["NUM", "ALPHA", "ALPHA_EXTENDED"]
            .into_iter()
            .map(String::from)
            .collect();
        assert_eq!(get_token_fast_classifier("ALLÉE", &names), None);
        assert_eq!(get_token_fast_classifier("RIVIÈRE", &names), None);
        assert_eq!(get_token_fast_classifier("SAINTE-THÉRÈSE", &names), None);
        // Plain ASCII still classifies through the (now single-pass) fast path.
        assert_eq!(
            get_token_fast_classifier("MAIN", &names),
            Some("ALPHA".to_string())
        );
    }
}