use anyhow::{bail, Result};
use burn::prelude::*;
use burn::store::{
    ModuleAdapter, ModuleSnapshot, PyTorchToBurnAdapter, SafetensorsStore, TensorSnapshot,
};
use burn::tensor::DType;
use std::path::Path;
use std::rc::Rc;

/// The weights ship as BF16, which the ndarray backend does not support.
/// burn's `HalfPrecisionAdapter` only covers F32<->F16, so BF16 -> F32 is
/// converted here, lazily.
#[derive(Debug, Clone, Default)]
struct Bf16ToF32;

impl ModuleAdapter for Bf16ToF32 {
    fn adapt(&self, snapshot: &TensorSnapshot) -> TensorSnapshot {
        if snapshot.dtype != DType::BF16 {
            return snapshot.clone();
        }
        let original = snapshot.clone_data_fn();
        let cast = Rc::new(move || Ok(original()?.convert_dtype(DType::F32)));
        TensorSnapshot::from_closure(
            cast,
            DType::F32,
            snapshot.shape.clone(),
            snapshot.path_stack.clone().unwrap_or_default(),
            snapshot.container_stack.clone().unwrap_or_default(),
            snapshot.tensor_id.unwrap_or_default(),
        )
    }

    fn clone_box(&self) -> Box<dyn ModuleAdapter> {
        Box::new(self.clone())
    }
}

use crate::model::Llama;

/// Renaming rules: Hugging Face names -> burn module paths.
/// Base Linear layers live under `.base.weight` because they are wrapped in LoraLinear.
const REMAP: &[(&str, &str)] = &[
    (r"^model\.embed_tokens\.weight$", "embed.weight"),
    (r"^model\.norm\.weight$", "norm.gamma"),
    (
        r"^model\.layers\.(\d+)\.input_layernorm\.weight$",
        "layers.${1}.attn_norm.gamma",
    ),
    (
        r"^model\.layers\.(\d+)\.post_attention_layernorm\.weight$",
        "layers.${1}.ffn_norm.gamma",
    ),
    (
        r"^model\.layers\.(\d+)\.self_attn\.q_proj\.weight$",
        "layers.${1}.attn.q.base.weight",
    ),
    (
        r"^model\.layers\.(\d+)\.self_attn\.k_proj\.weight$",
        "layers.${1}.attn.k.base.weight",
    ),
    (
        r"^model\.layers\.(\d+)\.self_attn\.v_proj\.weight$",
        "layers.${1}.attn.v.base.weight",
    ),
    (
        r"^model\.layers\.(\d+)\.self_attn\.o_proj\.weight$",
        "layers.${1}.attn.o.base.weight",
    ),
    (
        r"^model\.layers\.(\d+)\.mlp\.gate_proj\.weight$",
        "layers.${1}.ffn.gate.base.weight",
    ),
    (
        r"^model\.layers\.(\d+)\.mlp\.up_proj\.weight$",
        "layers.${1}.ffn.up.base.weight",
    ),
    (
        r"^model\.layers\.(\d+)\.mlp\.down_proj\.weight$",
        "layers.${1}.ffn.down.base.weight",
    ),
];

pub fn load_hf_weights<B: Backend>(model: &mut Llama<B>, dir: impl AsRef<Path>) -> Result<()> {
    let path = dir.as_ref().join("model.safetensors");
    if !path.exists() {
        bail!("{} does not exist", path.display());
    }

    let mut store = SafetensorsStore::from_file(&path)
        // BF16 -> F32, then transpose the Linear weights:
        // HF stores [out, in], burn expects [in, out]
        .with_from_adapter(Bf16ToF32.chain(PyTorchToBurnAdapter))
        // the LoRA tensors are not in the file; they keep their init
        .allow_partial(true);

    for (from, to) in REMAP {
        store = store.with_key_remapping(*from, *to);
    }

    let result = model.load_from(&mut store)?;

    println!("weights applied: {}", result.applied.len());
    if !result.missing.is_empty() {
        println!("missing        : {}", result.missing.len());
        for (k, _) in result.missing.iter().take(5) {
            println!("  - {k}");
        }
    }
    if !result.unused.is_empty() {
        println!("unused         : {}", result.unused.len());
        for k in result.unused.iter().take(5) {
            println!("  - {k}");
        }
    }
    Ok(())
}