//! safetensors loader: mmap, multi-shard (`model.safetensors.index.json`), bf16/f16/f32 -> f32, raw access for packed (NVFP4 / FP8) tensors.
//! Own ~100-line header parser instead of the `safetensors` crate so that raw byte slices of u8/fp8 tensors can be handed straight to CUDA upload.
use crate::tensor::Tensor;
use anyhow::{anyhow, bail, Context, Result};
use half::{bf16, f16};
use memmap2::Mmap;
use serde_json::Value;
use std::collections::BTreeMap;
use std::fs::File;
use std::path::{Path, PathBuf};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DType {
    F32,
    F16,
    BF16,
    F8E4M3,
    F8E5M2,
    U8,
    I8,
    I32,
    I64,
    Bool,
}

impl DType {
    fn parse(s: &str) -> Result<DType> {
        Ok(match s {
            "F32" => DType::F32,
            "F16" => DType::F16,
            "BF16" => DType::BF16,
            "F8_E4M3" => DType::F8E4M3,
            "F8_E5M2" => DType::F8E5M2,
            "U8" => DType::U8,
            "I8" => DType::I8,
            "I32" => DType::I32,
            "I64" => DType::I64,
            "BOOL" => DType::Bool,
            o => bail!("unsupported safetensors dtype {o}"),
        })
    }
    pub fn as_str(self) -> &'static str {
        match self {
            DType::F32 => "F32",
            DType::F16 => "F16",
            DType::BF16 => "BF16",
            DType::F8E4M3 => "F8_E4M3",
            DType::F8E5M2 => "F8_E5M2",
            DType::U8 => "U8",
            DType::I8 => "I8",
            DType::I32 => "I32",
            DType::I64 => "I64",
            DType::Bool => "BOOL",
        }
    }
    pub fn size(self) -> usize {
        match self {
            DType::F32 | DType::I32 => 4,
            DType::F16 | DType::BF16 => 2,
            DType::I64 => 8,
            _ => 1,
        }
    }
}

#[derive(Debug, Clone)]
pub struct TensorInfo {
    pub dtype: DType,
    pub shape: Vec<usize>,
    file: usize,
    start: usize,
    end: usize,
}

pub struct WeightStore {
    maps: Vec<Mmap>,
    index: BTreeMap<String, TensorInfo>,
}

impl WeightStore {
    /// `path` = a single .safetensors file, or a directory holding `model.safetensors` / `model-0000x-of-0000y.safetensors` shards.
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        let p = path.as_ref();
        let files: Vec<PathBuf> = if p.is_file() {
            vec![p.to_path_buf()]
        } else {
            let mut v: Vec<PathBuf> = std::fs::read_dir(p)
                .with_context(|| format!("read_dir {p:?}"))?
                .filter_map(|e| e.ok().map(|e| e.path()))
                .filter(|f| f.extension().map_or(false, |x| x == "safetensors"))
                .collect();
            v.sort();
            v
        };
        if files.is_empty() {
            bail!("no .safetensors in {p:?}");
        }
        let mut s = WeightStore { maps: vec![], index: BTreeMap::new() };
        for f in files {
            s.add_file(&f)?;
        }
        Ok(s)
    }

    fn add_file(&mut self, f: &Path) -> Result<()> {
        let file = File::open(f).with_context(|| format!("open {f:?}"))?;
        let mm = unsafe { Mmap::map(&file)? };
        if mm.len() < 8 {
            bail!("{f:?}: too short");
        }
        let n = u64::from_le_bytes(mm[..8].try_into().unwrap()) as usize;
        let hdr: Value = serde_json::from_slice(mm.get(8..8 + n).ok_or_else(|| anyhow!("{f:?}: bad header length"))?)?;
        let base = 8 + n;
        let fi = self.maps.len();
        for (k, v) in hdr.as_object().ok_or_else(|| anyhow!("header not object"))? {
            if k == "__metadata__" {
                continue;
            }
            let dtype = DType::parse(v["dtype"].as_str().unwrap_or(""))?;
            let shape: Vec<usize> = v["shape"].as_array().ok_or_else(|| anyhow!("shape"))?.iter().map(|x| x.as_u64().unwrap() as usize).collect();
            let off = v["data_offsets"].as_array().ok_or_else(|| anyhow!("offsets"))?;
            let (s, e) = (off[0].as_u64().unwrap() as usize, off[1].as_u64().unwrap() as usize);
            if base + e > mm.len() || e - s != shape.iter().product::<usize>() * dtype.size() {
                bail!("{f:?}: tensor {k} has inconsistent size");
            }
            self.index.insert(k.clone(), TensorInfo { dtype, shape, file: fi, start: base + s, end: base + e });
        }
        self.maps.push(mm);
        Ok(())
    }

    pub fn names(&self) -> impl Iterator<Item = &String> {
        self.index.keys()
    }
    pub fn contains(&self, name: &str) -> bool {
        self.index.contains_key(name)
    }
    pub fn info(&self, name: &str) -> Result<&TensorInfo> {
        self.index.get(name).ok_or_else(|| anyhow!("tensor not found: {name}"))
    }
    pub fn raw(&self, name: &str) -> Result<&[u8]> {
        let i = self.info(name)?;
        Ok(&self.maps[i.file][i.start..i.end])
    }
    /// any float dtype -> f32 tensor
    pub fn load_f32(&self, name: &str) -> Result<Tensor> {
        let i = self.info(name)?;
        let b = self.raw(name)?;
        let data: Vec<f32> = match i.dtype {
            DType::F32 => b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(),
            DType::BF16 => b.chunks_exact(2).map(|c| bf16::from_le_bytes(c.try_into().unwrap()).to_f32()).collect(),
            DType::F16 => b.chunks_exact(2).map(|c| f16::from_le_bytes(c.try_into().unwrap()).to_f32()).collect(),
            DType::F8E4M3 => b.iter().map(|&x| crate::quant::fp8_e4m3_to_f32(x)).collect(),
            d => bail!("{name}: cannot convert {d:?} to f32 directly"),
        };
        Ok(Tensor::new(i.shape.clone(), data))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;

    fn write_st(path: &Path) {
        // 2 tensors: a: f32 [2,2], b: bf16 [3]
        let a: Vec<u8> = [1.0f32, 2.0, 3.0, 4.5].iter().flat_map(|x| x.to_le_bytes()).collect();
        let b: Vec<u8> = [0.5f32, -1.0, 2.0].iter().flat_map(|x| bf16::from_f32(*x).to_le_bytes()).collect();
        let hdr = format!(
            r#"{{"__metadata__":{{"format":"pt"}},"a":{{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]}},"b":{{"dtype":"BF16","shape":[3],"data_offsets":[16,22]}}}}"#
        );
        let mut f = File::create(path).unwrap();
        f.write_all(&(hdr.len() as u64).to_le_bytes()).unwrap();
        f.write_all(hdr.as_bytes()).unwrap();
        f.write_all(&a).unwrap();
        f.write_all(&b).unwrap();
    }

    #[test]
    fn roundtrip_and_errors() {
        let dir = std::env::temp_dir().join(format!("dllm_loader_{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let p = dir.join("x.safetensors");
        write_st(&p);
        let s = WeightStore::open(&dir).unwrap();
        assert_eq!(s.names().collect::<Vec<_>>(), vec!["a", "b"]);
        let a = s.load_f32("a").unwrap();
        assert_eq!((a.shape.clone(), a.data.clone()), (vec![2, 2], vec![1.0, 2.0, 3.0, 4.5]));
        assert_eq!(s.load_f32("b").unwrap().data, vec![0.5, -1.0, 2.0]);
        assert!(s.load_f32("nope").is_err());
        // truncated file is rejected
        let bad = dir.join("bad.safetensors");
        std::fs::write(&bad, &std::fs::read(&p).unwrap()[..20]).unwrap();
        assert!(WeightStore::open(&bad).is_err());
        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn real_tiny_fixture_names() {
        let p = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../../ref/tiny/model.safetensors");
        if !p.exists() {
            return;
        }
        let s = WeightStore::open(&p).unwrap();
        assert_eq!(s.info("model.language_model.layers.2.self_attn.k_proj.weight").unwrap().shape, vec![32, 64]);
        assert!(!s.contains("model.language_model.layers.2.self_attn.v_proj.weight")); // k_eq_v full layer
        assert!(s.contains("model.language_model.layers.0.self_attn.v_proj.weight"));
    }
}