use crate::dto::InputKind;
use crate::errors::OrchestratorError;
use crate::prompt;
use std::path::Path;

pub struct TokenCounter {
    tok: tokenizers::Tokenizer,
}

impl TokenCounter {
    pub fn load(path: &Path) -> Result<Self, OrchestratorError> {
        let tok = tokenizers::Tokenizer::from_file(path)
            .map_err(|e| OrchestratorError::Tokenizer(e.to_string()))?;
        Ok(Self { tok })
    }

    /// Wrap the text per `input_type`, then tokenize with add_special_tokens=true.
    /// Returns the token count only.
    pub fn count_wrapped(
        &self,
        text: &str,
        kind: InputKind,
    ) -> Result<usize, OrchestratorError> {
        let wrapped = prompt::wrap(text, kind);
        let enc = self
            .tok
            .encode(wrapped, true)
            .map_err(|e| OrchestratorError::Tokenizer(e.to_string()))?;
        Ok(enc.get_ids().len())
    }
}

/// Round `n` up to the nearest multiple of `granularity`.
pub fn round_up(n: usize, granularity: usize) -> usize {
    if granularity == 0 {
        return n;
    }
    n.div_ceil(granularity) * granularity
}

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

    #[test]
    fn round_up_basics() {
        assert_eq!(round_up(8193, 64), 8256);
        assert_eq!(round_up(20000, 64), 20032);
        assert_eq!(round_up(1, 64), 64);
        assert_eq!(round_up(64, 64), 64);
    }
}