use tokenizers::Tokenizer;

pub struct TokenizerWrapper {
    tokenizer: Tokenizer,
}

impl TokenizerWrapper {
    pub fn from_file(path: &str) -> anyhow::Result<Self> {
        let tokenizer = Tokenizer::from_file(path)
            .map_err(|e| anyhow::anyhow!("加载 tokenizer 失败: {}", e))?;
        Ok(Self { tokenizer })
    }

    pub fn encode(&self, text: &str, add_special_tokens: bool) -> anyhow::Result<Vec<u32>> {
        let encoding = self
            .tokenizer
            .encode(text, add_special_tokens)
            .map_err(|e| anyhow::anyhow!("编码失败: {}", e))?;
        Ok(encoding.get_ids().to_vec())
    }

    pub fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> anyhow::Result<String> {
        let text = self
            .tokenizer
            .decode(ids, skip_special_tokens)
            .map_err(|e| anyhow::anyhow!("解码失败: {}", e))?;
        Ok(text)
    }

    pub fn vocab_size(&self) -> usize {
        self.tokenizer.get_vocab_size(true)
    }

    pub fn apply_chat_template(&self, user_msg: &str, enable_thinking: bool) -> anyhow::Result<Vec<u32>> {
        let prompt = if enable_thinking {
            format!(
                "<s><|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n<think>\n",
                user_msg
            )
        } else {
            format!(
                "<s><|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n",
                user_msg
            )
        };
        self.encode(&prompt, false)
    }
}