use burn::module::Module;
use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
use burn::tensor::backend::Backend;
use std::path::Path;

use crate::loader::{load_into_model, TensorStore};
use minicpm_core::config::LlamaConfig;
use minicpm_core::model::LlamaForCausalLM;

/// 导出全精度模型（f32）
pub fn export<B: Backend>(
    store: &TensorStore,
    config: &LlamaConfig,
    output_dir: &Path,
    device: &B::Device,
) -> anyhow::Result<()> {
    println!("创建模型结构...");
    let model = LlamaForCausalLM::<B>::new(config.clone(), device);

    println!("加载 safetensors 权重...");
    let model = load_into_model(model, store, device);

    println!("保存为 MPK 格式（全精度）...");
    std::fs::create_dir_all(output_dir)?;
    let output_path = output_dir.join("model");
    let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
    model.save_file(&output_path, &recorder)?;

    println!("全精度模型已导出到: {:?}", output_dir);
    Ok(())
}