pub mod config;
pub mod generator;
pub mod loader;
pub mod sampling;
pub mod tokenizer;

pub use config::GenerationConfig;
pub use generator::TokenStream;
pub use minicpm_core::config::{EosTokenId, LlamaConfig};
pub use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
pub use tokenizer::TokenizerWrapper;

use burn::tensor::backend::Backend;

use loader::{load_model, load_model_q8_from_dir};

pub struct MiniCPM<B: Backend> {
    model: LlamaForCausalLM<B>,
    config: LlamaConfig,
    tokenizer: TokenizerWrapper,
    device: B::Device,
}

impl<B: Backend> MiniCPM<B> {
    pub fn load(
        model_path: &str,
        config_path: &str,
        tokenizer_path: &str,
        device: &B::Device,
    ) -> anyhow::Result<Self> {
        let config = LlamaConfig::from_json(config_path)?;
        let tokenizer = TokenizerWrapper::from_file(tokenizer_path)?;
        let model = load_model::<B>(model_path, &config, device)?;

        Ok(Self {
            model,
            config,
            tokenizer,
            device: device.clone(),
        })
    }

    pub fn load_q8(
        model_dir: &str,
        config_path: &str,
        tokenizer_path: &str,
        device: &B::Device,
    ) -> anyhow::Result<Self> {
        let config = LlamaConfig::from_json(config_path)?;
        let tokenizer = TokenizerWrapper::from_file(tokenizer_path)?;
        let model = load_model_q8_from_dir::<B>(model_dir, &config, device)?;

        Ok(Self {
            model,
            config,
            tokenizer,
            device: device.clone(),
        })
    }

    pub fn generate(
        &self,
        prompt: &str,
        think: bool,
        config: &GenerationConfig,
    ) -> anyhow::Result<String> {
        let input_ids = self.tokenizer.apply_chat_template(prompt, think)?;
        let output_ids = generator::generate_tokens(
            &self.model,
            &input_ids,
            config,
            &self.config.eos_token_id,
            &self.device,
        );
        let new_ids = &output_ids[input_ids.len()..];
        self.tokenizer.decode(new_ids, true)
    }

    pub fn generate_stream<'a>(
        &'a self,
        prompt: &str,
        think: bool,
        config: &'a GenerationConfig,
    ) -> anyhow::Result<TextStream<'a, B>> {
        let input_ids = self.tokenizer.apply_chat_template(prompt, think)?;
        let token_stream = TokenStream::new(
            &self.model,
            &input_ids,
            config,
            &self.config.eos_token_id,
            &self.device,
        );

        Ok(TextStream {
            token_stream,
            tokenizer: &self.tokenizer,
            buffer: Vec::new(),
        })
    }

    pub fn config(&self) -> &LlamaConfig {
        &self.config
    }

    pub fn tokenizer(&self) -> &TokenizerWrapper {
        &self.tokenizer
    }

    pub fn inner_model(&self) -> &LlamaForCausalLM<B> {
        &self.model
    }
}

/// 文本流式迭代器 —— 解码 token 并返回文本片段
///
/// 用法示例：
/// ```ignore
/// let mut stream = model.generate_stream("你好", false, &config)?;
/// for text in stream {
///     print!("{}", text);
/// }
/// ```
pub struct TextStream<'a, B: Backend> {
    token_stream: TokenStream<'a, B>,
    tokenizer: &'a TokenizerWrapper,
    buffer: Vec<u32>,
}

impl<'a, B: Backend> Iterator for TextStream<'a, B> {
    type Item = String;

    fn next(&mut self) -> Option<Self::Item> {
        match self.token_stream.next() {
            Some(token) => {
                self.buffer.push(token);
                match self.tokenizer.decode(&[token], true) {
                    Ok(text) => Some(text),
                    Err(_) => Some(String::new()),
                }
            }
            None => None,
        }
    }
}

impl<'a, B: Backend> TextStream<'a, B> {
    pub fn all_tokens(&self) -> &[u32] {
        &self.buffer
    }
}