pub mod format;
pub mod loader;
pub mod utils;

pub use format::{ExportTask, Format};
pub use loader::TensorStore;

use burn::tensor::backend::Backend;
use std::path::Path;

use minicpm_core::config::LlamaConfig;

/// 执行导出任务
pub fn run_export<B: Backend>(
    safetensors_path: &Path,
    config_path: &Path,
    tokenizer_path: &Path,
    tasks: &[ExportTask],
    device: &B::Device,
) -> anyhow::Result<()> {
    println!("开始转换 MiniCPM 模型...");
    println!("源文件: {:?}", safetensors_path);

    println!("加载配置文件...");
    let config = LlamaConfig::from_json(config_path.to_str().unwrap())?;

    println!("解析 safetensors 文件...");
    let store = TensorStore::from_file(safetensors_path)?;

    for task in tasks {
        println!("\n--- 导出 {} ---", task.format.name());
        let output_dir = Path::new(&task.output_dir);

        match task.format {
            Format::Full => {
                format::full::export::<B>(&store, &config, output_dir, device)?;
            }
        }

        std::fs::copy(config_path, output_dir.join("config.json"))?;
        std::fs::copy(tokenizer_path, output_dir.join("tokenizer.json"))?;
    }

    println!("\n所有导出任务完成！");
    Ok(())
}