toklen 0.2.0

use serde::Deserialize;

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

use super::Error;

/// Raw deserialization helper for [`Metaspace`].
#[derive(Deserialize)]
struct MetaspaceRaw {
    #[serde(default = "String::new")]
    replacement: String,
    #[serde(default)]
    prepend_scheme: Option<String>,
    #[serde(default)]
    add_prefix_space: bool,
}

/// A compiled Metaspace pre-tokenizer.
///
/// Replaces whitespace characters with a replacement character (default `▁`)
/// and splits on the replacement. Widely used by Llama, T5, and other models.
#[derive(Clone, Debug, Deserialize)]
#[serde(try_from = "MetaspaceRaw")]
pub struct Metaspace {
    replacement: char,
    prepend_scheme: PrependScheme,
    add_prefix_space: bool,
}

#[derive(Clone, Debug, PartialEq, Eq)]
enum PrependScheme {
    Never,
    Always,
    First,
}

impl Metaspace {
    /// 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 mut new_buf = String::with_capacity(buffer.len());
        let mut new_splits = Vec::with_capacity(old_splits.len() << 1);

        let prepend_first = matches!(
            self.prepend_scheme,
            PrependScheme::Always | PrependScheme::First
        );
        let mut is_first_text = prepend_first;

        for split in old_splits {
            if split.token_id.is_some() {
                // Added token: encode through without modification.
                let start = new_buf.len();
                new_buf.push_str(&buffer[split.range.clone()]);
                let end = new_buf.len();
                new_splits.push(PtSplit {
                    range: start..end,
                    token_id: split.token_id,
                });
                continue;
            }

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

            let mut prev_was_space = false;

            for (_i, ch) in text.char_indices() {
                if ch.is_whitespace() {
                    prev_was_space = true;
                    continue;
                }

                // Start of a new segment.
                if prev_was_space || (is_first_text && self.add_prefix_space) {
                    let start = new_buf.len();
                    new_buf.push(self.replacement);
                    new_buf.push(ch);
                    let end = new_buf.len();
                    new_splits.push(PtSplit {
                        range: start..end,
                        token_id: None,
                    });
                    is_first_text = false;
                    prev_was_space = false;
                    continue;
                }

                // Non-space character within an existing segment: extend it.
                let c_start = new_buf.len();
                new_buf.push(ch);
                let c_end = new_buf.len();

                // Check if this continues the previous split or starts a new one.
                if let Some(last) = new_splits.last_mut()
                    && last.token_id.is_none()
                {
                    last.range.end = c_end;
                } else {
                    new_splits.push(PtSplit {
                        range: c_start..c_end,
                        token_id: None,
                    });
                }

                is_first_text = false;
                prev_was_space = false;
            }

            // If the text ends with whitespace, add a trailing replacement.
            if prev_was_space {
                let start = new_buf.len();
                new_buf.push(self.replacement);
                new_splits.push(PtSplit {
                    range: start..new_buf.len(),
                    token_id: None,
                });
            }
        }

        pts.buffer = new_buf;
        pts.splits = new_splits;

        Ok(())
    }
}

impl TryFrom<MetaspaceRaw> for Metaspace {
    type Error = Error;

    fn try_from(raw: MetaspaceRaw) -> Result<Self, Error> {
        let replacement = if raw.replacement.is_empty() {
            '\u{2581}' // Default: ▁ (LOWER ONE EIGHTH BLOCK)
        } else {
            raw.replacement
                .chars()
                .next()
                .ok_or_else(|| Error::Json(serde::de::Error::custom("empty replacement")))?
        };

        let prepend_scheme = match raw.prepend_scheme.as_deref() {
            None | Some("first") => PrependScheme::First,
            Some("always") => PrependScheme::Always,
            Some("never") => PrependScheme::Never,
            Some(other) => {
                return Err(Error::Unsupported(format!(
                    "unknown prepend_scheme: {other}"
                )));
            }
        };

        Ok(Self {
            replacement,
            prepend_scheme,
            add_prefix_space: raw.add_prefix_space,
        })
    }
}