//! The GGUF form (`extensions/prior-carriage.md` §9.2): version 3, little-endian — a header
//! of key-value metadata and tensor descriptions, then the tensors' bytes, each at a
//! multiple of the file's alignment.

use super::{refuse, DType, Refusal, Tensor, TensorSet};

/// The four bytes a GGUF file begins with.
pub const MAGIC: [u8; 4] = *b"GGUF";
/// The version this reader reads.
pub const VERSION: u32 = 3;
/// The alignment when `general.alignment` is absent.
pub const DEFAULT_ALIGNMENT: u64 = 32;
/// The longest tensor name, in bytes.
pub const MAX_NAME: usize = 64;
/// The most dimensions a tensor has.
pub const MAX_DIMS: usize = 4;
/// The metadata value type of a string.
pub const TYPE_STRING: u32 = 8;
const TYPE_ARRAY: u32 = 9;

struct Cur<'a> {
    b: &'a [u8],
    p: usize,
}

impl<'a> Cur<'a> {
    fn take(&mut self, n: u64) -> Result<&'a [u8], Refusal> {
        let n = usize::try_from(n).map_err(|_| Refusal::new("header_truncated", "a length past the address space"))?;
        let end = self.p.checked_add(n).filter(|&e| e <= self.b.len());
        let Some(end) = end else { return refuse("header_truncated", format!("{n} bytes due at byte {} of {}", self.p, self.b.len())) };
        let s = &self.b[self.p..end];
        self.p = end;
        Ok(s)
    }
    fn u32(&mut self) -> Result<u32, Refusal> {
        Ok(u32::from_le_bytes(self.take(4)?.try_into().expect("4")))
    }
    fn u64(&mut self) -> Result<u64, Refusal> {
        Ok(u64::from_le_bytes(self.take(8)?.try_into().expect("8")))
    }
    fn string(&mut self) -> Result<String, Refusal> {
        let n = self.u64()?;
        let s = self.take(n)?;
        String::from_utf8(s.to_vec()).map_err(|_| Refusal::new("header_malformed", "a string that is not UTF-8"))
    }
}

/// The size of a fixed-size metadata value type, or `None` for string and array.
fn fixed_size(t: u32) -> Option<u64> {
    match t {
        0 | 1 | 7 => Some(1),
        2 | 3 => Some(2),
        4..=6 => Some(4),
        10..=12 => Some(8),
        _ => None,
    }
}

/// Read one metadata value of type `t`; returns the string when it is one.
fn value(c: &mut Cur<'_>, t: u32) -> Result<Option<String>, Refusal> {
    if let Some(n) = fixed_size(t) {
        let v = c.take(n)?;
        if t == 7 && v[0] > 1 {
            return refuse("header_malformed", "a boolean that is neither 0 nor 1");
        }
        return Ok(None);
    }
    match t {
        TYPE_STRING => Ok(Some(c.string()?)),
        TYPE_ARRAY => {
            let et = c.u32()?;
            if et == TYPE_ARRAY {
                return refuse("header_malformed", "an array of arrays (this reader reads one level of array)");
            }
            if et != TYPE_STRING && fixed_size(et).is_none() {
                return refuse("header_malformed", format!("an array of value type {et}"));
            }
            let n = c.u64()?;
            // Each element takes at least one byte (a string at least eight), so a count
            // past the bytes left is a truncated header, refused before any reservation.
            let min = fixed_size(et).unwrap_or(8);
            if n.checked_mul(min).is_none_or(|need| need > (c.b.len() - c.p) as u64) {
                return refuse("header_truncated", format!("an array of {n} elements past the end"));
            }
            for _ in 0..n {
                value(c, et)?;
            }
            Ok(None)
        }
        _ => refuse("header_malformed", format!("a metadata value type {t}")),
    }
}

fn key_ok(k: &str) -> bool {
    !k.is_empty() && k.len() <= 65535 && k.split('.').all(|s| !s.is_empty() && s.bytes().all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'_'))
}

/// Read a GGUF file (§9.2): every rule of the format, no canonical check. String metadata
/// goes to the set's metadata; a value of another type is read and checked but not kept,
/// so a file that holds one is not canonical (§11.2).
pub fn read(b: &[u8]) -> Result<TensorSet, Refusal> {
    if b.len() < 4 {
        return refuse("header_truncated", "fewer than the 4 bytes of the magic");
    }
    if b[..4] != MAGIC {
        return refuse("bad_magic", "the file does not begin with GGUF");
    }
    let mut c = Cur { b, p: 4 };
    let version = c.u32()?;
    if version != VERSION {
        return refuse("version_unsupported", format!("version {version} (this reader reads version 3, little-endian)"));
    }
    let n_tensors = c.u64()?;
    let n_kv = c.u64()?;
    // A key-value pair takes at least 12 bytes, a tensor description at least 28.
    if n_kv.checked_mul(12).is_none_or(|v| v > b.len() as u64) || n_tensors.checked_mul(28).is_none_or(|v| v > b.len() as u64) {
        return refuse("header_truncated", format!("{n_kv} key-value pairs and {n_tensors} tensors past the end"));
    }
    let mut set = TensorSet::default();
    let mut keys = std::collections::BTreeSet::new();
    let mut alignment = DEFAULT_ALIGNMENT;
    for _ in 0..n_kv {
        let k = c.string()?;
        if !key_ok(&k) {
            return refuse("header_malformed", format!("key {k:?} is not lower_snake_case segments joined by '.'"));
        }
        if !keys.insert(k.clone()) {
            return refuse("header_malformed", format!("key {k:?} twice"));
        }
        let t = c.u32()?;
        let at = c.p;
        let s = value(&mut c, t)?;
        if k == "general.alignment" {
            let v = if t == 4 { u32::from_le_bytes(b[at..at + 4].try_into().expect("4")) as u64 } else { 0 };
            if t != 4 || v == 0 || v % 8 != 0 {
                return refuse("header_malformed", "general.alignment is not a non-zero uint32 multiple of 8");
            }
            alignment = v;
        }
        if let Some(s) = s {
            set.metadata.insert(k, s);
        }
    }
    struct Info {
        offset: u64,
        len: u64,
    }
    let mut infos = Vec::new();
    let mut names = std::collections::BTreeSet::new();
    for _ in 0..n_tensors {
        let name = c.string()?;
        if name.is_empty() || name.len() > MAX_NAME {
            return refuse("tensor_malformed", format!("tensor name {name:?} is empty or longer than {MAX_NAME} bytes"));
        }
        if !names.insert(name.clone()) {
            return refuse("tensor_malformed", format!("tensor {name:?} twice"));
        }
        let nd = c.u32()? as usize;
        if !(1..=MAX_DIMS).contains(&nd) {
            return refuse("tensor_malformed", format!("{name:?}: {nd} dimensions"));
        }
        let mut dims = Vec::with_capacity(nd);
        for _ in 0..nd {
            dims.push(c.u64()?);
        }
        if dims.contains(&0) {
            return refuse("tensor_malformed", format!("{name:?}: a dimension of 0"));
        }
        let ty = c.u32()?;
        let Some(dtype) = DType::from_gguf_type(ty) else { return refuse("dtype_unknown", format!("{name:?}: type {ty}")) };
        let offset = c.u64()?;
        if !offset.is_multiple_of(alignment) {
            return refuse("offset_mismatch", format!("{name:?}: offset {offset} is not a multiple of {alignment}"));
        }
        if dtype == DType::Q8_0 && dims[0] % 32 != 0 {
            return refuse("tensor_malformed", format!("{name:?}: Q8_0 rows of {} elements", dims[0]));
        }
        let count = dims.iter().try_fold(1u64, |a, &d| a.checked_mul(d));
        let Some(len) = count.and_then(|n| dtype.bytes_for(n)) else {
            return refuse("tensor_malformed", format!("{name:?}: the element count overflows"));
        };
        let shape: Vec<u64> = dims.iter().rev().copied().collect();
        infos.push(Info { offset, len });
        set.tensors.push(Tensor { name, dtype, shape, data: Vec::new() });
    }
    let header_end = c.p as u64;
    let data_start = header_end.div_ceil(alignment) * alignment;
    if data_start > b.len() as u64 {
        return refuse("header_truncated", "the file ends before its tensor data begins");
    }
    if b[header_end as usize..data_start as usize].iter().any(|&x| x != 0) {
        return refuse("header_malformed", "the padding after the header is not zero");
    }
    let data = &b[data_start as usize..];
    let mut order: Vec<usize> = (0..infos.len()).collect();
    order.sort_by_key(|&i| infos[i].offset);
    let mut at = 0u64;
    for &i in &order {
        let (o, l) = (infos[i].offset, infos[i].len);
        if o < at {
            return refuse("offset_mismatch", format!("{:?} overlaps the tensor before it", set.tensors[i].name));
        }
        let end = o.checked_add(l).filter(|&e| e <= data.len() as u64);
        let Some(end) = end else { return refuse("offset_mismatch", format!("{:?} runs past the end of the file", set.tensors[i].name)) };
        set.tensors[i].data = data[o as usize..end as usize].to_vec();
        at = end;
    }
    Ok(set.sorted())
}

/// The canonical bytes of `set` (§11.2): string metadata in ascending byte order of the
/// keys, the tensors in ascending byte order of their names, alignment 32 left implicit,
/// each tensor at the first multiple of 32 after the one before, zero padding, and nothing
/// after the last tensor.
pub fn write(set: &TensorSet) -> Result<Vec<u8>, Refusal> {
    let set = set.clone().sorted();
    let mut o = MAGIC.to_vec();
    o.extend_from_slice(&VERSION.to_le_bytes());
    o.extend_from_slice(&(set.tensors.len() as u64).to_le_bytes());
    o.extend_from_slice(&(set.metadata.len() as u64).to_le_bytes());
    let s = |o: &mut Vec<u8>, v: &str| {
        o.extend_from_slice(&(v.len() as u64).to_le_bytes());
        o.extend_from_slice(v.as_bytes());
    };
    for (k, v) in &set.metadata {
        if !key_ok(k) || k == "general.alignment" {
            return refuse("header_malformed", format!("key {k:?} cannot be written as a string"));
        }
        s(&mut o, k);
        o.extend_from_slice(&TYPE_STRING.to_le_bytes());
        s(&mut o, v);
    }
    let mut offsets = Vec::with_capacity(set.tensors.len());
    let mut at = 0u64;
    for (i, t) in set.tensors.iter().enumerate() {
        if i > 0 && set.tensors[i - 1].name == t.name {
            return refuse("tensor_malformed", format!("{:?} named twice", t.name));
        }
        if t.name.is_empty() || t.name.len() > MAX_NAME || !(1..=MAX_DIMS).contains(&t.shape.len()) || t.shape.contains(&0) {
            return refuse("tensor_malformed", format!("{:?} cannot be a GGUF tensor (name of 1 to 64 bytes, 1 to 4 dimensions, none 0)", t.name));
        }
        let Some(ty) = t.dtype.gguf_type() else { return refuse("dtype_unknown", format!("{:?}: {:?} has no GGUF type", t.name, t.dtype)) };
        if t.count().and_then(|n| t.dtype.bytes_for(n)) != Some(t.data.len() as u64) {
            return refuse("offset_mismatch", format!("{:?}: {} bytes for its shape and dtype", t.name, t.data.len()));
        }
        at = at.div_ceil(DEFAULT_ALIGNMENT) * DEFAULT_ALIGNMENT;
        offsets.push(at);
        s(&mut o, &t.name);
        o.extend_from_slice(&(t.shape.len() as u32).to_le_bytes());
        for d in t.shape.iter().rev() {
            o.extend_from_slice(&d.to_le_bytes());
        }
        o.extend_from_slice(&ty.to_le_bytes());
        o.extend_from_slice(&at.to_le_bytes());
        at += t.data.len() as u64;
    }
    let pad = |o: &mut Vec<u8>, base: usize| {
        while !((o.len() - base) as u64).is_multiple_of(DEFAULT_ALIGNMENT) {
            o.push(0);
        }
    };
    pad(&mut o, 0);
    let base = o.len();
    for (t, &off) in set.tensors.iter().zip(&offsets) {
        o.resize(base + off as usize, 0);
        o.extend_from_slice(&t.data);
    }
    Ok(o)
}

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

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

    #[test]
    fn writes_canonically_and_reads_back() {
        let bytes = write(&sample()).unwrap();
        assert_eq!(&bytes[..4], b"GGUF");
        let back = read(&bytes).unwrap();
        assert_eq!(back, sample().sorted());
        assert_eq!(write(&back).unwrap(), bytes);
        // Shape [2, 3] is written as dimensions [3, 2], innermost first.
        assert!(bytes.windows(9).any(|w| w == b"w\x02\x00\x00\x00\x03\x00\x00\x00"));
    }

    #[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(b"GG"), Some("header_truncated"));
        let mut m = good.clone();
        m[0] = b'X';
        assert_eq!(code(&m), Some("bad_magic"));
        let mut v = good.clone();
        v[4..8].copy_from_slice(&2u32.to_le_bytes());
        assert_eq!(code(&v), Some("version_unsupported"));
        let mut cnt = good.clone();
        cnt[8..16].copy_from_slice(&(u64::MAX / 2).to_le_bytes());
        assert_eq!(code(&cnt), Some("header_truncated"));
        assert_eq!(code(&good[..good.len() - 1]), Some("offset_mismatch"));
        // Trailing bytes read, but are not canonical.
        let mut t = good.clone();
        t.push(0);
        assert_ne!(write(&read(&t).unwrap()).unwrap(), t);
    }
}
