use charabia::Tokenize;
use std::collections::HashSet;
use tantivy::tokenizer::{Token, TokenStream, Tokenizer};

pub fn tokenize(text: &str, min_token_len: usize, max_token_len: usize) -> Vec<String> {
    let mut token_set = HashSet::new();
    text.tokenize()
        .filter(|token| {
            let lemma = token.lemma().to_string();
            lemma.len() >= min_token_len
                && lemma.len() <= max_token_len
                && token.kind() == charabia::TokenKind::Word
                && token_set.insert(lemma)
        })
        .map(|token| token.lemma().to_string())
        .collect()
}

#[derive(Clone, Debug)]
pub struct CharabiaTokenizer {
    min_token_len: usize,
    max_token_len: usize,
    token: Token,
}

impl CharabiaTokenizer {
    pub fn new(min_token_len: usize, max_token_len: usize) -> Self {
        CharabiaTokenizer {
            min_token_len,
            max_token_len,
            token: Token::default(),
        }
    }

    pub fn name() -> &'static str {
        "charabia"
    }
}

pub struct CharabiaTokenStream<'a> {
    token: &'a mut Token,
    tokens: Vec<String>,
    index: usize,
}

impl TokenStream for CharabiaTokenStream<'_> {
    fn advance(&mut self) -> bool {
        if let Some(token) = self.tokens.get(self.index) {
            self.token.text.clear();
            self.token.text.push_str(token);
            self.token.position = 0;
            self.index += 1;
            true
        } else {
            false
        }
    }

    fn token(&self) -> &Token {
        self.token
    }

    fn token_mut(&mut self) -> &mut Token {
        self.token
    }
}

impl Tokenizer for CharabiaTokenizer {
    type TokenStream<'a> = CharabiaTokenStream<'a>;

    fn token_stream<'a>(&'a mut self, text: &'a str) -> CharabiaTokenStream<'a> {
        self.token.reset();
        CharabiaTokenStream {
            token: &mut self.token,
            tokens: tokenize(text, self.min_token_len, self.max_token_len),
            index: 0,
        }
    }
}