toklen 0.2.0

use serde::Deserialize;

use crate::pre_tokenized::PreTokenizedString;
use crate::pre_tokenized::PtSplit;

use super::Error;

/// How punctuation is handled in the output.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub enum Behavior {
    #[default]
    Isolated,
    Removed,
    MergedWithPrevious,
    MergedWithNext,
    Contiguous,
}

/// A compiled Punctuation pre-tokenizer.
///
/// Isolates or removes punctuation from surrounding text, matching the
/// HuggingFace `Punctuation` pre-tokenizer.
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct Punctuation {
    #[serde(default)]
    behavior: Behavior,
}

impl Punctuation {
    /// Refine the splits of a [`PreTokenizedString`] in place.
    pub fn pre_tokenize(&self, pts: &mut PreTokenizedString) -> Result<(), Error> {
        let buffer = pts.buffer.as_str();
        let old_splits = &pts.splits;
        let hint: usize = old_splits.len() << 1; // Vec::len <= isize::MAX
        let mut new_splits = Vec::with_capacity(hint);

        for split in old_splits {
            if split.token_id.is_some() {
                new_splits.push(split.clone());
                continue;
            }

            let text = &buffer[split.range.clone()];
            if text.is_empty() {
                continue;
            }

            let base = split.range.start;
            let segments = find_punctuation_segments(text);
            let ranges = apply_behavior(self.behavior, &segments);

            for (s, e) in ranges {
                if s < e {
                    new_splits.push(PtSplit {
                        range: (base + s)..(base + e),
                        token_id: None,
                    });
                }
            }
        }

        pts.splits = new_splits;
        Ok(())
    }
}

/// Segment the input into alternating punctuation / non-punctuation spans.
fn find_punctuation_segments(input: &str) -> Vec<(usize, usize, bool)> {
    let mut segments = Vec::new();
    let mut current_start = 0;
    let mut current_is_punct = is_punctuation_char(input.chars().next());

    for (i, ch) in input.char_indices() {
        let is_punct = is_punctuation_char(Some(ch));
        if is_punct != current_is_punct {
            if i > current_start {
                segments.push((current_start, i, current_is_punct));
            }
            current_start = i;
            current_is_punct = is_punct;
        }
    }
    if current_start < input.len() {
        segments.push((current_start, input.len(), current_is_punct));
    }

    segments
}

/// Check if a character is punctuation (Unicode category P).
fn is_punctuation_char(c: Option<char>) -> bool {
    match c {
        Some(ch) => {
            // Fast path: ASCII punctuation.
            if ch.is_ascii_punctuation() {
                return true;
            }
            // Unicode punctuation: check if not alphanumeric, not whitespace,
            // not a combining mark, and not a symbol-math/currency.
            // This is a pragmatic approximation of Unicode category P.
            if ch.is_alphabetic() || ch.is_numeric() || ch.is_whitespace() {
                return false;
            }
            // The remainder includes punctuation, symbols, marks, etc.
            // For tokenizer purposes, treat anything that isn't letter,
            // number, or whitespace as "punctuation."
            true
        }
        None => false,
    }
}

/// Apply punctuation behavior to segments, returning final (start, end) pairs.
fn apply_behavior(behavior: Behavior, segments: &[(usize, usize, bool)]) -> Vec<(usize, usize)> {
    match behavior {
        Behavior::Removed => segments
            .iter()
            .filter(|&&(_, _, is_punct)| !is_punct)
            .map(|&(s, e, _)| (s, e))
            .collect(),

        Behavior::Isolated => segments.iter().map(|&(s, e, _)| (s, e)).collect(),

        Behavior::Contiguous => {
            let mut result: Vec<(usize, usize)> = Vec::new();
            let mut prev_is_punct = None;
            for &(s, e, is_punct) in segments {
                if prev_is_punct == Some(is_punct) {
                    if let Some(last) = result.last_mut() {
                        last.1 = e;
                    }
                } else {
                    result.push((s, e));
                }
                prev_is_punct = Some(is_punct);
            }
            result
        }

        Behavior::MergedWithPrevious => {
            let mut result: Vec<(usize, usize)> = Vec::new();
            let mut prev_was_punct = false;
            for &(s, e, is_punct) in segments {
                if is_punct && !prev_was_punct {
                    if let Some(last) = result.last_mut() {
                        last.1 = e;
                    } else {
                        result.push((s, e));
                    }
                } else {
                    result.push((s, e));
                }
                prev_was_punct = is_punct;
            }
            result
        }

        Behavior::MergedWithNext => {
            let mut result: Vec<(usize, usize)> = Vec::new();
            let mut prev_was_punct = false;
            for &(s, e, is_punct) in segments.iter().rev() {
                if is_punct && !prev_was_punct {
                    if let Some(last) = result.last_mut() {
                        last.0 = s;
                    } else {
                        result.push((s, e));
                    }
                } else {
                    result.push((s, e));
                }
                prev_was_punct = is_punct;
            }
            result.reverse();
            result
        }
    }
}