use anyhow::{Result, bail};
use once_cell::sync::Lazy;

/// Supported LLM token encodings.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum TokenEncoding {
    /// cl100k_base — used by GPT-3.5, GPT-4, text-embedding-ada-002.
    #[default]
    Cl100kBase,
    /// o200k_base — used by GPT-4o, GPT-5, o1, o3, o4 models.
    O200kBase,
}

impl std::str::FromStr for TokenEncoding {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self> {
        match s.to_lowercase().as_str() {
            "cl100k_base" => Ok(TokenEncoding::Cl100kBase),
            "o200k_base" => Ok(TokenEncoding::O200kBase),
            _ => bail!("Unknown token encoding: {s}. Supported: cl100k_base, o200k_base"),
        }
    }
}

impl std::fmt::Display for TokenEncoding {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TokenEncoding::Cl100kBase => write!(f, "cl100k_base"),
            TokenEncoding::O200kBase => write!(f, "o200k_base"),
        }
    }
}

/// Wrapper around tiktoken BPE tokenizer.
///
/// Provides streaming-friendly tokenization: count tokens in text,
/// split text into token strings, and decode token IDs back to text.
#[derive(Clone)]
pub struct Tokenizer {
    bpe: tiktoken_rs::CoreBPE,
}

impl std::fmt::Debug for Tokenizer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Tokenizer").finish()
    }
}

/// Static tokenizer instances — loaded once per process, shared across all plugins.
static CL100K: Lazy<Tokenizer> = Lazy::new(|| {
    Tokenizer::new(TokenEncoding::Cl100kBase).expect("Failed to create cl100k_base tokenizer")
});
static O200K: Lazy<Tokenizer> = Lazy::new(|| {
    Tokenizer::new(TokenEncoding::O200kBase).expect("Failed to create o200k_base tokenizer")
});

/// Returns a reference to a cached tokenizer for the given encoding.
///
/// The BPE vocabulary is loaded once per encoding and reused for the
/// lifetime of the process.
pub fn get_tokenizer(encoding: TokenEncoding) -> &'static Tokenizer {
    match encoding {
        TokenEncoding::Cl100kBase => &CL100K,
        TokenEncoding::O200kBase => &O200K,
    }
}

impl Tokenizer {
    /// Creates a new tokenizer for the specified encoding.
    pub fn new(encoding: TokenEncoding) -> Result<Self> {
        let bpe = match encoding {
            TokenEncoding::Cl100kBase => tiktoken_rs::cl100k_base()?,
            TokenEncoding::O200kBase => tiktoken_rs::o200k_base()?,
        };
        Ok(Self { bpe })
    }

    /// Counts the number of tokens in the given text.
    ///
    /// Uses `encode_ordinary` which treats the text as a single unit.
    /// For streaming: tokenizing chunks independently and summing gives
    /// the same result as tokenizing the full text (exact when no regex
    /// match spans a chunk boundary).
    pub fn count(&self, text: &str) -> usize {
        self.bpe.encode_ordinary(text).len()
    }

    /// Splits text into individual decoded token strings.
    ///
    /// Each returned string corresponds to one token. Useful for finding
    /// exact byte boundaries when filtering by token count.
    pub fn split_by_token(&self, text: &str) -> Result<Vec<String>> {
        self.bpe.split_by_token(text, false)
    }

    /// Returns an iterator over decoded token strings.
    ///
    /// Lazily produces token strings without allocating a Vec for all tokens.
    /// Use this when you only need the first N tokens (e.g., head/skip filters).
    pub fn split_by_token_iter<'a>(
        &'a self,
        text: &'a str,
    ) -> impl Iterator<Item = Result<String>> + 'a {
        self.bpe.split_by_token_iter(text, false)
    }

    /// Counts tokens up to `max_tokens` and returns `(token_count, byte_position)`.
    ///
    /// Uses an iterator to stop early, avoiding allocation of token strings
    /// beyond `max_tokens`. The byte_position is in the lossy UTF-8 encoding
    /// of `text` — use `map_lossy_pos_to_bytes` to map back to original bytes.
    pub fn count_bounded(&self, text: &str, max_tokens: usize) -> (usize, usize) {
        let mut count = 0usize;
        let mut byte_pos = 0usize;
        for token_str in self.bpe.split_by_token_iter(text, false) {
            let Ok(s) = token_str else {
                break;
            };
            byte_pos += s.len();
            count += 1;
            if count >= max_tokens {
                break;
            }
        }
        (count, byte_pos)
    }

    /// Decodes a slice of token IDs back into a string.
    pub fn decode(&self, tokens: &[u32]) -> Result<String> {
        self.bpe.decode(tokens.to_vec())
    }
}

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

    #[test]
    fn test_tokenizer_count() {
        let tok = Tokenizer::new(TokenEncoding::Cl100kBase).unwrap();
        let count = tok.count("Hello, world!");
        assert!(count > 0);
    }

    #[test]
    fn test_tokenizer_split() {
        let tok = Tokenizer::new(TokenEncoding::Cl100kBase).unwrap();
        let tokens = tok.split_by_token("Hello world").unwrap();
        assert_eq!(tokens.len(), 2);
    }

    #[test]
    fn test_tokenizer_roundtrip() {
        let tok = Tokenizer::new(TokenEncoding::Cl100kBase).unwrap();
        let text = "The quick brown fox jumps over the lazy dog.";
        let token_ids: Vec<u32> = tok
            .bpe
            .encode_ordinary(text)
            .into_iter()
            .map(|x| x as u32)
            .collect();
        let decoded = tok.decode(&token_ids).unwrap();
        assert_eq!(text, decoded);
    }

    #[test]
    fn test_chunk_sum_close_to_full() {
        let tok = Tokenizer::new(TokenEncoding::Cl100kBase).unwrap();
        let text = "The quick brown fox jumps over the lazy dog. \
                     Pack my box with five dozen liquor jugs. \
                     How vexingly quick daft zebras jump!";
        let full_count = tok.count(text);

        // Split into chunks at word boundaries
        let mid = text.find("Pack").unwrap();
        let (a, b) = text.split_at(mid);
        let chunk_sum = tok.count(a) + tok.count(b);
        // Chunk-based counting may differ by 1-2 tokens when a BPE merge
        // boundary falls near the chunk split point
        assert!(
            (full_count as isize - chunk_sum as isize).abs() <= 2,
            "full={full_count}, chunk_sum={chunk_sum}"
        );
    }

    #[test]
    fn test_encoding_from_str() {
        assert_eq!(
            "cl100k_base".parse::<TokenEncoding>().unwrap(),
            TokenEncoding::Cl100kBase
        );
        assert_eq!(
            "o200k_base".parse::<TokenEncoding>().unwrap(),
            TokenEncoding::O200kBase
        );
        assert!("unknown".parse::<TokenEncoding>().is_err());
    }
}
