//! The safetensors form (`extensions/prior-carriage.md` §9.1): an 8-byte little-endian
//! header length `N`, `N` bytes of JSON header, then the tensors' bytes.

use super::{refuse, DType, Refusal, Tensor, TensorSet};
use crate::strict_json::{self, quote, Json};

/// The largest header the reader accepts, in bytes (§9.1 step 2).
pub const MAX_HEADER: u64 = 100_000_000;
/// The member of the header that holds the metadata.
pub const METADATA: &str = "__metadata__";

/// Read a safetensors file (§9.1): every rule of the format, no canonical check.
pub fn read(b: &[u8]) -> Result<TensorSet, Refusal> {
    if b.len() < 8 {
        return refuse("header_truncated", "fewer than the 8 bytes of the header length");
    }
    let n = u64::from_le_bytes(b[..8].try_into().expect("8 bytes"));
    if n > MAX_HEADER {
        return refuse("header_too_large", format!("a header of {n} bytes, past {MAX_HEADER}"));
    }
    let end = 8 + n as usize;
    if end > b.len() {
        return refuse("header_truncated", format!("a header of {n} bytes in a file of {} bytes", b.len()));
    }
    let doc = strict_json::read(&b[8..end]).map_err(|e| Refusal::new("header_malformed", e))?;
    let Json::Obj(members) = doc else { return refuse("header_malformed", "the header is not a JSON object") };
    let data = &b[end..];
    let mut set = TensorSet::default();
    let mut spans: Vec<(u64, u64, usize)> = Vec::new();
    for (name, v) in members {
        if name == METADATA {
            let Json::Obj(m) = v else { return refuse("header_malformed", "__metadata__ is not an object") };
            for (k, mv) in m {
                let Json::Str(s) = mv else { return refuse("header_malformed", format!("metadata {k:?} is not a string")) };
                set.metadata.insert(k, s);
            }
            continue;
        }
        if name.is_empty() {
            return refuse("tensor_malformed", "a tensor with an empty name");
        }
        let Json::Obj(t) = &v else { return refuse("tensor_malformed", format!("{name:?} is not an object")) };
        if t.len() != 3 || ["dtype", "shape", "data_offsets"].iter().any(|k| v.get(k).is_none()) {
            return refuse("tensor_malformed", format!("{name:?} does not have exactly dtype, shape and data_offsets"));
        }
        let Some(ds) = v.get("dtype").and_then(Json::as_str) else {
            return refuse("tensor_malformed", format!("{name:?}: dtype is not a string"));
        };
        let Some(dtype) = DType::from_safetensors_name(ds) else {
            return refuse("dtype_unknown", format!("{name:?}: dtype {ds:?}"));
        };
        let shape: Option<Vec<u64>> = v.get("shape").and_then(Json::as_arr).and_then(|a| a.iter().map(Json::as_u64).collect());
        let Some(shape) = shape else { return refuse("tensor_malformed", format!("{name:?}: shape is not a list of integers")) };
        let offs: Option<Vec<u64>> = v.get("data_offsets").and_then(Json::as_arr).and_then(|a| a.iter().map(Json::as_u64).collect());
        let Some(offs) = offs.filter(|o| o.len() == 2 && o[0] <= o[1]) else {
            return refuse("tensor_malformed", format!("{name:?}: data_offsets is not [begin, end] with begin ≤ end"));
        };
        let count = shape.iter().try_fold(1u64, |a, &d| a.checked_mul(d));
        let Some(need) = count.and_then(|c| dtype.bytes_for(c)) else {
            return refuse("tensor_malformed", format!("{name:?}: the element count overflows"));
        };
        if offs[1] - offs[0] != need {
            return refuse("offset_mismatch", format!("{name:?}: data_offsets span {} bytes, its shape and dtype {need}", offs[1] - offs[0]));
        }
        spans.push((offs[0], offs[1], set.tensors.len()));
        set.tensors.push(Tensor { name, dtype, shape, data: Vec::new() });
    }
    // The spans cover the data exactly: no gap, no overlap, nothing after the last.
    spans.sort();
    let mut at = 0u64;
    for &(b0, e0, i) in &spans {
        if b0 != at {
            return refuse("offset_mismatch", format!("{:?} begins at {b0}, the data so far ends at {at}", set.tensors[i].name));
        }
        at = e0;
    }
    if at != data.len() as u64 {
        return refuse("offset_mismatch", format!("the tensors cover {at} bytes of {} bytes of data", data.len()));
    }
    for &(b0, e0, i) in &spans {
        set.tensors[i].data = data[b0 as usize..e0 as usize].to_vec();
    }
    Ok(set.sorted())
}

/// The canonical header text of `set` (§11.1), before its padding.
pub fn header(set: &TensorSet) -> Result<String, Refusal> {
    let mut parts: Vec<String> = Vec::new();
    if !set.metadata.is_empty() {
        let m: Vec<String> = set.metadata.iter().map(|(k, v)| format!("{}:{}", quote(k), quote(v))).collect();
        parts.push(format!("{}:{{{}}}", quote(METADATA), m.join(",")));
    }
    let mut at = 0u64;
    for t in &set.tensors {
        let Some(ds) = t.dtype.safetensors_name() else {
            return refuse("dtype_unknown", format!("{:?}: {:?} has no safetensors name", t.name, t.dtype));
        };
        let end = at + t.data.len() as u64;
        let shape: Vec<String> = t.shape.iter().map(u64::to_string).collect();
        parts.push(format!("{}:{{\"dtype\":\"{ds}\",\"shape\":[{}],\"data_offsets\":[{at},{end}]}}", quote(&t.name), shape.join(",")));
        at = end;
    }
    Ok(format!("{{{}}}", parts.join(",")))
}

/// The canonical bytes of `set` (§11.1): the tensors in ascending byte order of their
/// names, the header padded with spaces so that the data begins at a multiple of 8.
pub fn write(set: &TensorSet) -> Result<Vec<u8>, Refusal> {
    let set = set.clone().sorted();
    for w in set.tensors.windows(2) {
        if w[0].name == w[1].name {
            return refuse("tensor_malformed", format!("{:?} named twice", w[0].name));
        }
    }
    for t in &set.tensors {
        if t.name.is_empty() || t.name == METADATA {
            return refuse("tensor_malformed", format!("{:?} is not a tensor name", t.name));
        }
        let need = t.count().and_then(|c| t.dtype.bytes_for(c));
        if need != Some(t.data.len() as u64) {
            return refuse("offset_mismatch", format!("{:?}: {} bytes for its shape and dtype", t.name, t.data.len()));
        }
    }
    let mut h = header(&set)?.into_bytes();
    while (8 + h.len()) % 8 != 0 {
        h.push(b' ');
    }
    let mut out = (h.len() as u64).to_le_bytes().to_vec();
    out.extend_from_slice(&h);
    for t in &set.tensors {
        out.extend_from_slice(&t.data);
    }
    Ok(out)
}

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

    fn sample() -> TensorSet {
        let mut s = TensorSet::default();
        s.metadata.insert("wai.profile".into(), "wai-int/1".into());
        s.tensors.push(Tensor::from_ints("b", DType::I32, vec![2], &[-1, 7]));
        s.tensors.push(Tensor::from_ints("a", DType::I8, vec![1, 3], &[1, -2, 3]));
        s
    }

    #[test]
    fn writes_canonically_and_reads_back() {
        let bytes = write(&sample()).unwrap();
        let n = u64::from_le_bytes(bytes[..8].try_into().unwrap()) as usize;
        assert_eq!((8 + n) % 8, 0);
        assert_eq!(
            std::str::from_utf8(&bytes[8..8 + n]).unwrap().trim_end(),
            r#"{"__metadata__":{"wai.profile":"wai-int/1"},"a":{"dtype":"I8","shape":[1,3],"data_offsets":[0,3]},"b":{"dtype":"I32","shape":[2],"data_offsets":[3,11]}}"#
        );
        let back = read(&bytes).unwrap();
        assert_eq!(back, sample().sorted());
        assert_eq!(write(&back).unwrap(), bytes);
    }

    #[test]
    fn refuses_what_the_format_forbids() {
        let good = write(&sample()).unwrap();
        let code = |b: &[u8]| read(b).err().map(|r| r.code);
        assert_eq!(code(&good[..5]), Some("header_truncated"));
        let mut big = good.clone();
        big[..8].copy_from_slice(&(MAX_HEADER + 1).to_le_bytes());
        assert_eq!(code(&big), Some("header_too_large"));
        let with_header = |h: &str, data: &[u8]| {
            let mut o = (h.len() as u64).to_le_bytes().to_vec();
            o.extend_from_slice(h.as_bytes());
            o.extend_from_slice(data);
            o
        };
        assert_eq!(code(&with_header(r#"{"a":{"dtype":"I8","shape":[1],"data_offsets":[0,1]},"a":{}}"#, &[0])), Some("header_malformed"));
        assert_eq!(code(&with_header(r#"{"a":{"dtype":"Q4","shape":[1],"data_offsets":[0,1]}}"#, &[0])), Some("dtype_unknown"));
        assert_eq!(code(&with_header(r#"{"a":{"dtype":"I8","shape":[2],"data_offsets":[0,1]}}"#, &[0])), Some("offset_mismatch"));
        assert_eq!(code(&with_header(r#"{"a":{"dtype":"I8","shape":[1],"data_offsets":[0,1]}}"#, &[0, 0])), Some("offset_mismatch"));
        assert_eq!(code(&with_header(r#"{"a":{"dtype":"I8","shape":[1],"data_offsets":[0,1],"x":1}}"#, &[0])), Some("tensor_malformed"));
        assert_eq!(code(&with_header(r#"{"__metadata__":{"k":1}}"#, &[])), Some("header_malformed"));
        // Reads, and is not canonical: a space before the object.
        let loose = with_header(r#" {"a":{"dtype":"I8","shape":[1],"data_offsets":[0,1]}}"#, &[5]);
        let set = read(&loose).unwrap();
        assert_ne!(write(&set).unwrap(), loose);
    }
}
