use anyhow::{bail, Result};
use burn::prelude::*;
use burn::store::{ModuleSnapshot, SafetensorsStore};
use std::path::Path;

use crate::config::LoraConfig;
use crate::model::Llama;

const WEIGHTS: &str = "adapter.safetensors";
const CONFIG: &str = "adapter_config.json";

/// Only the adapter tensors: `....lora.a` and `....lora.b`.
/// Everything else (base weights, norms, embeddings) stays out of the file.
const LORA_RE: &str = r"\.lora\.[ab]$";

/// Saves only the adapter. No `ModuleAdapter` is applied: these tensors are
/// created and loaded in burn layout, so transposing them like the HF
/// weights would corrupt them.
pub fn save<B: Backend>(model: &Llama<B>, lora: &LoraConfig, dir: impl AsRef<Path>) -> Result<()> {
    let dir = dir.as_ref();
    std::fs::create_dir_all(dir)?;

    let mut store = SafetensorsStore::from_file(dir.join(WEIGHTS)).with_regex(LORA_RE);
    model.save_into(&mut store)?;

    std::fs::write(dir.join(CONFIG), serde_json::to_string_pretty(lora)?)?;
    Ok(())
}

/// Reads the adapter config. It is needed BEFORE building the model: the
/// rank defines the shape of A and B, and the targets which projections get them.
pub fn load_config(dir: impl AsRef<Path>) -> Result<LoraConfig> {
    let path = dir.as_ref().join(CONFIG);
    if !path.exists() {
        bail!("{} does not exist", path.display());
    }
    Ok(serde_json::from_str(&std::fs::read_to_string(&path)?)?)
}

/// Applies an adapter to a model already loaded with the base weights.
pub fn load<B: Backend>(model: &mut Llama<B>, dir: impl AsRef<Path>) -> Result<usize> {
    let path = dir.as_ref().join(WEIGHTS);
    if !path.exists() {
        bail!("{} does not exist", path.display());
    }

    // `allow_partial`: the file only holds LoRA; the rest of the model keeps
    // the base weights that were already loaded.
    let mut store = SafetensorsStore::from_file(&path).allow_partial(true);
    let result = model.load_from(&mut store)?;

    if !result.errors.is_empty() {
        bail!("errors applying the adapter: {:?}", result.errors);
    }
    // `allow_partial` also lets an incomplete file through: tensors that are
    // not applied would keep their random init and silently corrupt the
    // output. Require all of them.
    let expected = model.trainable_ids().len();
    if result.applied.len() != expected {
        bail!(
            "the adapter applied {} of {} expected tensors; different rank or targets?",
            result.applied.len(),
            expected
        );
    }
    Ok(result.applied.len())
}