bbpe 0.6.3

//! Configuration builders controlling training and corpus ingestion.

use std::collections::HashSet;
use std::convert::TryFrom;

use crate::error::{BbpeError, Result};
use crate::special_tokens;
use serde::{Deserialize, Serialize};

fn default_true() -> bool {
    true
}

/// Preprocessing mode applied before BPE training.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct PreprocessorConfig {
    /// Which splitter to apply ahead of training.
    pub kind: PreprocessorKind,
    /// Probability that a detected delimiter boundary is preserved (1.0 = always split).
    pub split_probability: f64,
    /// Optional RNG seed for deterministic probabilistic preprocessing.
    pub seed: Option<u64>,
}

impl Default for PreprocessorConfig {
    fn default() -> Self {
        Self {
            kind: PreprocessorKind::None,
            split_probability: 1.0,
            seed: None,
        }
    }
}

impl PreprocessorConfig {
    /// Validates probability bounds and returns an error when misconfigured.
    pub fn validate(&self) -> Result<()> {
        if !(0.0..=1.0).contains(&self.split_probability) {
            return Err(BbpeError::InvalidConfig(format!(
                "preprocessor split probability ({}) must be between 0.0 and 1.0",
                self.split_probability
            )));
        }
        Ok(())
    }
}

/// Supported preprocessing strategies.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum PreprocessorKind {
    /// Disable preprocessing, operating directly on raw byte streams.
    None,
    /// Split on contiguous ASCII whitespace runs (space, tab, CR/LF, vertical tab, form feed).
    AsciiWhitespace,
    /// Split on Unicode whitespace using [`char::is_whitespace`].
    UnicodeWhitespace,
    /// Split binary sequences on runs of `0x00` bytes.
    NullDelimited,
}

/// Configuration for binary BPE training.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct TrainerConfig {
    /// Target vocabulary size including the 256 base byte tokens and special tokens.
    pub target_vocab_size: usize,
    /// Minimum number of pair occurrences required before a merge is considered.
    pub min_frequency: usize,
    /// Inclusive list of allowed token lengths (in bytes) produced by merges.
    pub allowed_token_lengths: Vec<usize>,
    /// Enables per-iteration logging through the `log` facade.
    pub show_progress: bool,
    /// Additional tokens appended to the vocabulary after training.
    pub special_tokens: Vec<String>,
    /// Frequency threshold below which merges are considered part of a plateau.
    pub plateau_frequency_floor: usize,
    /// Number of consecutive plateau iterations before considering early stopping.
    pub plateau_patience: usize,
    /// Ratio between the initial and current pair frequency that also signals a plateau.
    pub plateau_frequency_divisor: usize,
    /// Hard cap on merge iterations; `None` uses the target vocabulary size.
    pub max_merge_iterations: Option<usize>,
    /// Enables plateau-based early stopping with `plateau_patience`.
    pub plateau_stop_enabled: bool,
    /// Optional preprocessing performed before counting merges.
    pub preprocessor: PreprocessorConfig,
    /// Require ASCII letters when merges produce whitespace-containing tokens.
    #[serde(default)]
    pub require_letter_whitespace_merges: bool,
    /// Forbid merges whose tokens begin with ASCII whitespace unless the token is pure whitespace.
    #[serde(default)]
    pub forbid_leading_whitespace_merges: bool,
    /// Controls whether reasoning/argument tokens are inserted after the byte alphabet.
    #[serde(default = "default_true")]
    pub reasoning_tokens_enabled: bool,
}

impl TrainerConfig {
    /// Returns a builder initialised with [`TrainerConfig::default`].
    #[must_use]
    pub fn builder() -> TrainerBuilder {
        TrainerBuilder::default()
    }

    /// Validates the invariants required for training.
    pub fn validate(&self) -> Result<()> {
        self.preprocessor.validate()?;
        let leading_specials = special_tokens::leading_tokens().len();
        let reasoning_specials = if self.reasoning_tokens_enabled {
            special_tokens::reasoning_tokens().len()
        } else {
            0
        };
        let min_vocab = leading_specials + 256 + reasoning_specials + self.special_tokens.len();
        if self.target_vocab_size < min_vocab {
            return Err(BbpeError::InvalidConfig(format!(
                "target_vocab_size ({}) must be at least {} (leading specials + 256 byte tokens + trailing specials {}).",
                self.target_vocab_size,
                min_vocab,
                self.special_tokens.len()
            )));
        }
        if self.min_frequency == 0 {
            return Err(BbpeError::InvalidConfig(
                "min_frequency must be greater than zero".into(),
            ));
        }
        let max_vocab = usize::try_from(u32::MAX).unwrap_or(usize::MAX);
        if self.target_vocab_size > max_vocab {
            return Err(BbpeError::InvalidConfig(format!(
                "target_vocab_size ({}) exceeds {max_vocab}, the maximum representable TokenId",
                self.target_vocab_size
            )));
        }
        if !self.allowed_token_lengths.contains(&1) {
            return Err(BbpeError::InvalidConfig(
                "allowed_token_lengths must include the base length of 1".into(),
            ));
        }
        if self.plateau_frequency_divisor == 0 {
            return Err(BbpeError::InvalidConfig(
                "plateau_frequency_divisor must be greater than zero".into(),
            ));
        }
        if self.plateau_stop_enabled && self.plateau_patience == 0 {
            return Err(BbpeError::InvalidConfig(
                "plateau_patience must be > 0 when plateau_stop_enabled is true".into(),
            ));
        }
        if self.allowed_token_lengths.is_empty() {
            return Err(BbpeError::InvalidConfig(
                "allowed_token_lengths must not be empty".into(),
            ));
        }
        Ok(())
    }
}

impl Default for TrainerConfig {
    fn default() -> Self {
        Self {
            target_vocab_size: 32_768,
            min_frequency: 4,
            allowed_token_lengths: (1..=32).collect(),
            show_progress: true,
            special_tokens: Vec::new(),
            plateau_frequency_floor: 128,
            plateau_patience: 32,
            plateau_frequency_divisor: 512,
            max_merge_iterations: None,
            plateau_stop_enabled: false,
            preprocessor: PreprocessorConfig::default(),
            require_letter_whitespace_merges: false,
            forbid_leading_whitespace_merges: false,
            reasoning_tokens_enabled: true,
        }
    }
}

/// Builder for [`TrainerConfig`].
#[derive(Debug, Clone)]
pub struct TrainerBuilder {
    cfg: TrainerConfig,
    allowed_lengths_overridden: bool,
    special_tokens_overridden: bool,
    appended_special_tokens: Vec<String>,
}

impl Default for TrainerBuilder {
    fn default() -> Self {
        let mut cfg = TrainerConfig::default();
        cfg.special_tokens.clear();
        Self {
            cfg,
            allowed_lengths_overridden: false,
            special_tokens_overridden: false,
            appended_special_tokens: Vec::new(),
        }
    }
}

impl TrainerBuilder {
    /// Creates a builder with [`TrainerConfig::default`] settings.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets the desired vocabulary size (including base byte tokens).
    #[must_use]
    pub fn target_vocab_size(mut self, value: usize) -> Self {
        self.cfg.target_vocab_size = value;
        self
    }

    /// Sets the minimum merge frequency.
    #[must_use]
    pub fn min_frequency(mut self, value: usize) -> Self {
        self.cfg.min_frequency = value;
        self
    }

    /// Overrides the allowed token lengths.
    #[must_use]
    pub fn allowed_token_lengths<I>(mut self, lengths: I) -> Self
    where
        I: IntoIterator<Item = usize>,
    {
        self.cfg.allowed_token_lengths = lengths.into_iter().collect();
        self.allowed_lengths_overridden = true;
        self
    }

    /// Enables or disables per-iteration logging.
    #[must_use]
    pub fn show_progress(mut self, enabled: bool) -> Self {
        self.cfg.show_progress = enabled;
        self
    }

    /// Appends custom special tokens after the built-in required tokens.
    #[must_use]
    pub fn special_tokens<I, S>(mut self, tokens: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.special_tokens_overridden = true;
        self.cfg.special_tokens = tokens.into_iter().map(|s| s.into()).collect();
        self
    }

    /// Appends additional special tokens after the reasoning set (preserves order, deduplicated in build).
    #[must_use]
    pub fn append_special_tokens<I, S>(mut self, tokens: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.appended_special_tokens
            .extend(tokens.into_iter().map(|s| s.into()));
        self
    }

    /// Configures plateau-based early stopping thresholds.
    #[must_use]
    pub fn plateau_thresholds(mut self, floor: usize, patience: usize, divisor: usize) -> Self {
        self.cfg.plateau_frequency_floor = floor;
        self.cfg.plateau_patience = patience;
        self.cfg.plateau_frequency_divisor = divisor;
        self
    }

    /// Enables plateau-based early stopping.
    #[must_use]
    pub fn plateau_stop_enabled(mut self, enabled: bool) -> Self {
        self.cfg.plateau_stop_enabled = enabled;
        self
    }

    /// Sets a hard merge iteration limit.
    #[must_use]
    pub fn max_merge_iterations(mut self, value: Option<usize>) -> Self {
        self.cfg.max_merge_iterations = value;
        self
    }

    /// Configures the preprocessing mode applied before training.
    #[must_use]
    pub fn preprocessor(mut self, config: PreprocessorConfig) -> Self {
        self.cfg.preprocessor = config;
        self
    }

    /// Enforces that whitespace-containing merges include ASCII letters when enabled.
    #[must_use]
    pub fn require_letter_whitespace_merges(mut self, enabled: bool) -> Self {
        self.cfg.require_letter_whitespace_merges = enabled;
        self
    }

    /// Forbids merges that would introduce leading ASCII whitespace unless the token is pure whitespace.
    #[must_use]
    pub fn forbid_leading_whitespace_merges(mut self, enabled: bool) -> Self {
        self.cfg.forbid_leading_whitespace_merges = enabled;
        self
    }

    /// Enables or disables the optional reasoning/argument special tokens.
    #[must_use]
    pub fn reasoning_tokens_enabled(mut self, enabled: bool) -> Self {
        self.cfg.reasoning_tokens_enabled = enabled;
        self
    }

    /// Overrides only the preprocessor split probability.
    #[must_use]
    pub fn preprocessor_split_probability(mut self, probability: f64) -> Self {
        self.cfg.preprocessor.split_probability = probability;
        self
    }

    /// Sets an optional RNG seed for probabilistic preprocessing.
    #[must_use]
    pub fn preprocessor_seed(mut self, seed: Option<u64>) -> Self {
        self.cfg.preprocessor.seed = seed;
        self
    }

    /// Finalises the builder, returning a validated [`TrainerConfig`].
    pub fn build(mut self) -> Result<TrainerConfig> {
        if !self.allowed_lengths_overridden {
            match self.cfg.preprocessor.kind {
                PreprocessorKind::AsciiWhitespace | PreprocessorKind::UnicodeWhitespace => {
                    self.cfg.allowed_token_lengths = (1..=16).collect();
                }
                _ => {}
            }
        }
        self.cfg.allowed_token_lengths.sort_unstable();
        self.cfg.allowed_token_lengths.dedup();
        let mut trailing = std::mem::take(&mut self.cfg.special_tokens);
        if self.special_tokens_overridden {
            trailing.extend(self.appended_special_tokens.iter().cloned());
        } else {
            let leading: HashSet<String> =
                special_tokens::leading_tokens().iter().cloned().collect();
            trailing.extend(
                self.appended_special_tokens
                    .iter()
                    .filter(|token| !leading.contains(token.as_str()))
                    .cloned(),
            );
        }
        special_tokens::dedup_in_place(&mut trailing);
        self.cfg.special_tokens = trailing;
        self.cfg.validate()?;
        Ok(self.cfg)
    }
}

/// Configuration controlling how binary corpora are read from disk.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct IngestConfig {
    /// Size of chunks to read from each input file; `0` reads entire files.
    pub chunk_size: usize,
    /// Enables recursive directory traversal.
    pub recursive: bool,
    /// Follows symlinks encountered during traversal.
    pub follow_symlinks: bool,
}

impl Default for IngestConfig {
    fn default() -> Self {
        Self {
            chunk_size: 8192,
            recursive: true,
            follow_symlinks: false,
        }
    }
}

impl IngestConfig {
    /// Returns a builder initialised with [`IngestConfig::default`].
    #[must_use]
    pub fn builder() -> IngestBuilder {
        IngestBuilder::default()
    }
}

/// Builder for [`IngestConfig`].
#[derive(Debug, Default, Clone)]
pub struct IngestBuilder {
    cfg: IngestConfig,
}

impl IngestBuilder {
    /// Creates a new builder with [`IngestConfig::default`] settings.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets the chunk size in bytes (0 = read entire file at once).
    #[must_use]
    pub fn chunk_size(mut self, size: usize) -> Self {
        self.cfg.chunk_size = size;
        self
    }

    /// Enables or disables recursive directory traversal.
    #[must_use]
    pub fn recursive(mut self, enabled: bool) -> Self {
        self.cfg.recursive = enabled;
        self
    }

    /// Enables or disables following of symlinks when traversing directories.
    #[must_use]
    pub fn follow_symlinks(mut self, enabled: bool) -> Self {
        self.cfg.follow_symlinks = enabled;
        self
    }

    /// Finalises the builder, returning the [`IngestConfig`].
    pub fn build(self) -> IngestConfig {
        self.cfg
    }
}

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

    #[test]
    fn builder_deduplicates_allowed_lengths() {
        let cfg = TrainerConfig::builder()
            .allowed_token_lengths([4, 2, 2, 1])
            .show_progress(false)
            .build()
            .expect("config should be valid");
        assert_eq!(&cfg.allowed_token_lengths, &[1, 2, 4]);
    }

    #[test]
    fn ascii_preprocessor_shrinks_default_allowed_lengths() {
        let cfg = TrainerConfig::builder()
            .preprocessor(PreprocessorConfig {
                kind: PreprocessorKind::AsciiWhitespace,
                split_probability: 1.0,
                seed: None,
            })
            .build()
            .expect("config build");
        assert_eq!(cfg.allowed_token_lengths.first(), Some(&1));
        assert_eq!(cfg.allowed_token_lengths.last(), Some(&16));
    }

    #[test]
    fn custom_allowed_lengths_are_preserved() {
        let cfg = TrainerConfig::builder()
            .allowed_token_lengths([1, 4, 24, 32])
            .preprocessor(PreprocessorConfig {
                kind: PreprocessorKind::AsciiWhitespace,
                split_probability: 1.0,
                seed: None,
            })
            .build()
            .expect("config build");
        assert_eq!(cfg.allowed_token_lengths, vec![1, 4, 24, 32]);
    }

    #[test]
    fn validate_rejects_missing_base_length() {
        let cfg = TrainerConfig {
            allowed_token_lengths: vec![2, 3],
            ..TrainerConfig::default()
        };
        let err = cfg.validate().expect_err("validation should fail");
        assert!(matches!(
            err,
            BbpeError::InvalidConfig(message) if message.contains("allowed_token_lengths must include")
        ));
    }

    #[test]
    fn ingest_builder_overrides_defaults() {
        let cfg = IngestConfig::builder()
            .chunk_size(1024)
            .recursive(false)
            .follow_symlinks(true)
            .build();
        assert_eq!(cfg.chunk_size, 1024);
        assert!(!cfg.recursive);
        assert!(cfg.follow_symlinks);
    }
}