pub use minicpm_core::config::{EosTokenId, LlamaConfig};
pub use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};

use burn::module::Module;
use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
use burn::tensor::backend::Backend;
use burn::tensor::{Float, Int, Tensor};
use rand::Rng;
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!("Failed to load 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!("Failed to encode: {}", 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!("Failed to decode: {}", e))?;
        Ok(text)
    }

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

    /// 应用 MiniCPM5 chat template
    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)
    }
}


pub struct GenerationConfig {
    pub max_new_tokens: Option<usize>,
    pub temperature: f32,
    pub top_p: f32,
}

impl Default for GenerationConfig {
    fn default() -> Self {
        Self {
            max_new_tokens: None,
            temperature: 1.0,
            top_p: 1.0,
        }
    }
}

pub fn generate_with_cache<B: Backend>(
    model: &LlamaForCausalLM<B>,
    input_ids: &[u32],
    config: &GenerationConfig,
    eos_token_id: &EosTokenId,
    device: &B::Device,
) -> Vec<u32> {
    let mut output_ids = input_ids.to_vec();
    let mut cache: Option<LlamaKVCache<B>> = None;

    // 第一步：输入完整 prompt，建立初始 cache
    {
        let input_ints: Vec<i64> = input_ids.iter().map(|&x| x as i64).collect();
        let input_tensor =
            Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), device)
                .unsqueeze::<2>();

        let (logits, new_cache) = model.forward_with_cache(input_tensor, cache.as_ref());
        cache = Some(new_cache);

        let next_token = sample_last(&logits, config.temperature, config.top_p);
        output_ids.push(next_token);

        if eos_token_id.contains(next_token) {
            return output_ids;
        }
    }

    // 后续步骤：每次只输入 1 个 token，使用 cache
    let mut count = 1;
    loop {
        if let Some(max) = config.max_new_tokens {
            if count >= max {
                break;
            }
        }

        let last_token = output_ids[output_ids.len() - 1];
        let input_ints: Vec<i64> = vec![last_token as i64];
        let input_tensor =
            Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), device)
                .unsqueeze::<2>();

        let (logits, new_cache) = model.forward_with_cache(input_tensor, cache.as_ref());
        cache = Some(new_cache);

        let vocab_size = model.config.vocab_size;
        let next_token_logits = logits
            .clone()
            .slice([0..1, 0..1, 0..vocab_size])
            .reshape([vocab_size]);

        let next_token = sample(&next_token_logits, config.temperature, config.top_p);
        output_ids.push(next_token);

        if eos_token_id.contains(next_token) {
            break;
        }

        count += 1;
    }

    output_ids
}

fn sample_last<B: Backend>(logits: &Tensor<B, 3, Float>, temperature: f32, top_p: f32) -> u32 {
    let shape = logits.shape();
    let dims: [usize; 3] = shape.dims();
    let seq_len = dims[1];
    let vocab_size = dims[2];

    let last_logits = logits
        .clone()
        .slice([0..1, seq_len - 1..seq_len, 0..vocab_size])
        .reshape([vocab_size]);

    sample(&last_logits, temperature, top_p)
}

fn sample<B: Backend>(logits: &Tensor<B, 1, Float>, temperature: f32, top_p: f32) -> u32 {
    let data = logits.clone().to_data();
    let values: Vec<f32> = data.as_slice().unwrap().to_vec();

    // temperature = 0 时退化为 greedy
    if temperature <= 0.001 {
        let mut max_idx = 0;
        let mut max_val = f32::NEG_INFINITY;
        for (i, &val) in values.iter().enumerate() {
            if val > max_val {
                max_val = val;
                max_idx = i;
            }
        }
        return max_idx as u32;
    }

    // 应用 temperature
    let scaled: Vec<f32> = values.iter().map(|&v| v / temperature).collect();

    // softmax
    let max_val = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let exp_sum: f32 = scaled.iter().map(|&v| (v - max_val).exp()).sum();
    let mut probs: Vec<f32> = scaled.iter().map(|&v| (v - max_val).exp() / exp_sum).collect();

    // top_p 过滤
    if top_p < 1.0 {
        let mut indexed: Vec<(usize, f32)> = probs.iter().enumerate().map(|(i, &p)| (i, p)).collect();
        indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

        let mut cumsum = 0.0;
        let mut cutoff = indexed.len();
        for (i, &(_, p)) in indexed.iter().enumerate() {
            cumsum += p;
            if cumsum > top_p {
                cutoff = i + 1;
                break;
            }
        }

        let keep: std::collections::HashSet<usize> = indexed[..cutoff].iter().map(|(i, _)| *i).collect();
        for (i, p) in probs.iter_mut().enumerate() {
            if !keep.contains(&i) {
                *p = 0.0;
            }
        }

        let sum: f32 = probs.iter().sum();
        if sum > 0.0 {
            for p in probs.iter_mut() {
                *p /= sum;
            }
        }
    }

    // 多项式采样
    let mut rng = rand::thread_rng();
    let r: f32 = rng.gen();
    let mut cumsum = 0.0;
    for (i, &p) in probs.iter().enumerate() {
        cumsum += p;
        if r < cumsum {
            return i as u32;
        }
    }
    probs.len() as u32 - 1
}


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 = LlamaForCausalLM::<B>::new(config.clone(), device);
        let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
        let model = model.load_file(model_path, &recorder, 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 = generate_with_cache(
            &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 config(&self) -> &LlamaConfig {
        &self.config
    }

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

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