modelshelf 0.1.0

//! Minimal GGUF header/metadata reader.
//!
//! Reads just enough of a GGUF file to extract identification metadata
//! (`general.name`, `general.architecture`, quantization, context length)
//! without depending on ggml and without loading tensor data. Supports GGUF
//! versions 1–3. Unknown value types abort metadata parsing gracefully:
//! whatever was extracted up to that point is returned.

use std::io::{BufReader, Read, Seek, SeekFrom};
use std::path::Path;

use crate::registry::GgufMeta;
use crate::{Error, Result};

/// GGUF magic bytes.
pub const MAGIC: &[u8; 4] = b"GGUF";

/// Sanity cap on the metadata kv count (real files have dozens).
const MAX_KV_COUNT: u64 = 1 << 20;
/// Sanity cap on a single string we actually materialize.
const MAX_STRING_LEN: u64 = 1 << 20;

/// Returns true if the file starts with the GGUF magic.
pub fn is_gguf(path: &Path) -> bool {
    let Ok(mut f) = std::fs::File::open(path) else {
        return false;
    };
    let mut magic = [0u8; 4];
    std::io::Read::read_exact(&mut f, &mut magic).is_ok() && &magic == MAGIC
}

/// Read identification metadata from a GGUF file.
pub fn read_meta(path: &Path) -> Result<GgufMeta> {
    let f = std::fs::File::open(path).map_err(|e| Error::io(path, e))?;
    let mut r = BufReader::new(f);
    parse_meta(&mut r).map_err(|e| match e {
        Error::Io { source, .. } => Error::io(path, source),
        other => other,
    })
}

fn parse_meta<R: Read + Seek>(r: &mut R) -> Result<GgufMeta> {
    let mut magic = [0u8; 4];
    read_exact(r, &mut magic)?;
    if &magic != MAGIC {
        return Err(Error::Refused("not a GGUF file (bad magic)".into()));
    }
    let version = read_u32(r)?;
    if !(1..=3).contains(&version) {
        return Err(Error::Refused(format!(
            "unsupported GGUF version {version}"
        )));
    }
    // v1 uses 32-bit counts and lengths; v2+ use 64-bit.
    let wide = version >= 2;
    let _tensor_count = read_count(r, wide)?;
    let kv_count = read_count(r, wide)?;
    if kv_count > MAX_KV_COUNT {
        return Err(Error::Refused(format!(
            "implausible GGUF metadata kv count {kv_count}"
        )));
    }

    let mut meta = GgufMeta::default();
    let mut file_type: Option<u64> = None;
    let mut architecture: Option<String> = None;
    let mut wanted_ctx: Option<u64> = None;

    for _ in 0..kv_count {
        let Some(key) = read_string(r, wide)? else {
            return Ok(finish(meta, architecture, file_type, wanted_ctx));
        };
        let value = match read_value(r, wide) {
            Ok(v) => v,
            // Unknown value type: we cannot know its length, so we cannot
            // skip past it. Return what we have.
            Err(_) => return Ok(finish(meta, architecture, file_type, wanted_ctx)),
        };
        match key.as_str() {
            "general.name" => {
                if let Value::Str(s) = value {
                    meta.general_name = Some(s);
                }
            }
            "general.architecture" => {
                if let Value::Str(s) = value {
                    architecture = Some(s);
                }
            }
            "general.file_type" => {
                if let Some(n) = value.as_u64() {
                    file_type = Some(n);
                }
            }
            "general.parameter_count" => {
                if let Some(n) = value.as_u64() {
                    meta.parameter_count = Some(n);
                }
            }
            _ => {
                // `<arch>.context_length`: remember any such key; the right
                // one is picked once the architecture is known (keys can
                // appear in any order, but in practice arch comes first).
                if key.ends_with(".context_length") {
                    if let Some(n) = value.as_u64() {
                        wanted_ctx = Some(n);
                    }
                }
            }
        }
    }
    Ok(finish(meta, architecture, file_type, wanted_ctx))
}

fn finish(
    mut meta: GgufMeta,
    architecture: Option<String>,
    file_type: Option<u64>,
    ctx: Option<u64>,
) -> GgufMeta {
    meta.architecture = architecture;
    meta.quantization = file_type.map(file_type_label);
    meta.context_length = ctx;
    meta
}

/// Map `general.file_type` (llama_ftype) to a human-readable label.
fn file_type_label(ft: u64) -> String {
    match ft {
        0 => "F32".into(),
        1 => "F16".into(),
        2 => "Q4_0".into(),
        3 => "Q4_1".into(),
        7 => "Q8_0".into(),
        8 => "Q5_0".into(),
        9 => "Q5_1".into(),
        10 => "Q2_K".into(),
        11 => "Q3_K_S".into(),
        12 => "Q3_K_M".into(),
        13 => "Q3_K_L".into(),
        14 => "Q4_K_S".into(),
        15 => "Q4_K_M".into(),
        16 => "Q5_K_S".into(),
        17 => "Q5_K_M".into(),
        18 => "Q6_K".into(),
        19 => "IQ2_XXS".into(),
        20 => "IQ2_XS".into(),
        21 => "Q2_K_S".into(),
        22 => "IQ3_XS".into(),
        23 => "IQ3_XXS".into(),
        24 => "IQ1_S".into(),
        25 => "IQ4_NL".into(),
        26 => "IQ3_S".into(),
        27 => "IQ3_M".into(),
        28 => "IQ2_S".into(),
        29 => "IQ2_M".into(),
        30 => "IQ4_XS".into(),
        31 => "IQ1_M".into(),
        32 => "BF16".into(),
        other => format!("unknown({other})"),
    }
}

// ---- low-level readers -------------------------------------------------

enum Value {
    Str(String),
    U64(u64),
    Other,
}

impl Value {
    fn as_u64(&self) -> Option<u64> {
        match self {
            Value::U64(n) => Some(*n),
            _ => None,
        }
    }
}

fn read_exact<R: Read>(r: &mut R, buf: &mut [u8]) -> Result<()> {
    r.read_exact(buf).map_err(|e| Error::io("<gguf>", e))
}

fn read_u32<R: Read>(r: &mut R) -> Result<u32> {
    let mut b = [0u8; 4];
    read_exact(r, &mut b)?;
    Ok(u32::from_le_bytes(b))
}

fn read_u64<R: Read>(r: &mut R) -> Result<u64> {
    let mut b = [0u8; 8];
    read_exact(r, &mut b)?;
    Ok(u64::from_le_bytes(b))
}

fn read_count<R: Read>(r: &mut R, wide: bool) -> Result<u64> {
    if wide {
        read_u64(r)
    } else {
        Ok(read_u32(r)? as u64)
    }
}

/// Read a length-prefixed UTF-8 string. Returns `None` (instead of erroring)
/// for implausibly long strings, signaling the caller to stop parsing.
fn read_string<R: Read + Seek>(r: &mut R, wide: bool) -> Result<Option<String>> {
    let len = read_count(r, wide)?;
    if len > MAX_STRING_LEN {
        return Ok(None);
    }
    let mut buf = vec![0u8; len as usize];
    read_exact(r, &mut buf)?;
    Ok(Some(String::from_utf8_lossy(&buf).into_owned()))
}

/// Skip a length-prefixed string without materializing it.
fn skip_string<R: Read + Seek>(r: &mut R, wide: bool) -> Result<()> {
    let len = read_count(r, wide)?;
    r.seek(SeekFrom::Current(len as i64))
        .map_err(|e| Error::io("<gguf>", e))?;
    Ok(())
}

/// GGUF metadata value type ids.
const T_UINT8: u32 = 0;
const T_INT8: u32 = 1;
const T_UINT16: u32 = 2;
const T_INT16: u32 = 3;
const T_UINT32: u32 = 4;
const T_INT32: u32 = 5;
const T_FLOAT32: u32 = 6;
const T_BOOL: u32 = 7;
const T_STRING: u32 = 8;
const T_ARRAY: u32 = 9;
const T_UINT64: u32 = 10;
const T_INT64: u32 = 11;
const T_FLOAT64: u32 = 12;

fn fixed_size(ty: u32) -> Option<u64> {
    match ty {
        T_UINT8 | T_INT8 | T_BOOL => Some(1),
        T_UINT16 | T_INT16 => Some(2),
        T_UINT32 | T_INT32 | T_FLOAT32 => Some(4),
        T_UINT64 | T_INT64 | T_FLOAT64 => Some(8),
        _ => None,
    }
}

fn read_value<R: Read + Seek>(r: &mut R, wide: bool) -> Result<Value> {
    let ty = read_u32(r)?;
    match ty {
        T_UINT8 | T_INT8 | T_BOOL => {
            let mut b = [0u8; 1];
            read_exact(r, &mut b)?;
            Ok(Value::U64(b[0] as u64))
        }
        T_UINT16 | T_INT16 => {
            let mut b = [0u8; 2];
            read_exact(r, &mut b)?;
            Ok(Value::U64(u16::from_le_bytes(b) as u64))
        }
        T_UINT32 | T_INT32 => Ok(Value::U64(read_u32(r)? as u64)),
        T_UINT64 | T_INT64 => Ok(Value::U64(read_u64(r)?)),
        T_FLOAT32 => {
            r.seek(SeekFrom::Current(4))
                .map_err(|e| Error::io("<gguf>", e))?;
            Ok(Value::Other)
        }
        T_FLOAT64 => {
            r.seek(SeekFrom::Current(8))
                .map_err(|e| Error::io("<gguf>", e))?;
            Ok(Value::Other)
        }
        T_STRING => Ok(match read_string(r, wide)? {
            Some(s) => Value::Str(s),
            None => Value::Other,
        }),
        T_ARRAY => {
            let elem_ty = read_u32(r)?;
            let count = read_count(r, wide)?;
            if let Some(sz) = fixed_size(elem_ty) {
                let bytes = count
                    .checked_mul(sz)
                    .ok_or_else(|| Error::Refused("GGUF array too large".into()))?;
                r.seek(SeekFrom::Current(bytes as i64))
                    .map_err(|e| Error::io("<gguf>", e))?;
            } else if elem_ty == T_STRING {
                // e.g. tokenizer vocabularies: skip element by element.
                for _ in 0..count {
                    skip_string(r, wide)?;
                }
            } else {
                return Err(Error::Refused(format!(
                    "unskippable GGUF array element type {elem_ty}"
                )));
            }
            Ok(Value::Other)
        }
        other => Err(Error::Refused(format!("unknown GGUF value type {other}"))),
    }
}

#[cfg(any(test, feature = "test-fixtures"))]
pub mod fixtures {
    //! Tiny valid-GGUF writer for tests. Not part of the public API.

    use std::io::Write;
    use std::path::Path;

    fn put_string(out: &mut Vec<u8>, s: &str) {
        out.extend_from_slice(&(s.len() as u64).to_le_bytes());
        out.extend_from_slice(s.as_bytes());
    }

    /// Serialize a minimal GGUF v3 file with identification metadata and an
    /// arbitrary payload appended (so fixtures of the same metadata can still
    /// have distinct content hashes).
    pub fn gguf_bytes(
        name: &str,
        arch: &str,
        file_type: u32,
        ctx_len: u32,
        payload: &[u8],
    ) -> Vec<u8> {
        let mut out = Vec::new();
        out.extend_from_slice(b"GGUF");
        out.extend_from_slice(&3u32.to_le_bytes()); // version
        out.extend_from_slice(&0u64.to_le_bytes()); // tensor count
        out.extend_from_slice(&4u64.to_le_bytes()); // kv count

        put_string(&mut out, "general.name");
        out.extend_from_slice(&8u32.to_le_bytes()); // string
        put_string(&mut out, name);

        put_string(&mut out, "general.architecture");
        out.extend_from_slice(&8u32.to_le_bytes());
        put_string(&mut out, arch);

        put_string(&mut out, "general.file_type");
        out.extend_from_slice(&4u32.to_le_bytes()); // uint32
        out.extend_from_slice(&file_type.to_le_bytes());

        put_string(&mut out, &format!("{arch}.context_length"));
        out.extend_from_slice(&4u32.to_le_bytes());
        out.extend_from_slice(&ctx_len.to_le_bytes());

        out.extend_from_slice(payload);
        out
    }

    /// Write a fixture GGUF file to `path`, creating parent directories.
    pub fn write_gguf(
        path: &Path,
        name: &str,
        arch: &str,
        file_type: u32,
        ctx_len: u32,
        payload: &[u8],
    ) {
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent).unwrap();
        }
        let mut f = std::fs::File::create(path).unwrap();
        f.write_all(&gguf_bytes(name, arch, file_type, ctx_len, payload))
            .unwrap();
    }
}

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

    #[test]
    fn parses_fixture_metadata() {
        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("model.gguf");
        fixtures::write_gguf(&path, "Test Llama 8B", "llama", 15, 8192, b"PAYLOAD");
        assert!(is_gguf(&path));
        let meta = read_meta(&path).unwrap();
        assert_eq!(meta.general_name.as_deref(), Some("Test Llama 8B"));
        assert_eq!(meta.architecture.as_deref(), Some("llama"));
        assert_eq!(meta.quantization.as_deref(), Some("Q4_K_M"));
        assert_eq!(meta.context_length, Some(8192));
    }

    #[test]
    fn rejects_non_gguf() {
        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("not-a-model.bin");
        std::fs::write(&path, b"hello world, definitely not gguf").unwrap();
        assert!(!is_gguf(&path));
        assert!(read_meta(&path).is_err());
    }

    #[test]
    fn survives_string_arrays_like_tokenizer_vocab() {
        // Hand-build a GGUF with a big-ish string array before general.name.
        let mut out: Vec<u8> = Vec::new();
        out.extend_from_slice(b"GGUF");
        out.extend_from_slice(&3u32.to_le_bytes());
        out.extend_from_slice(&0u64.to_le_bytes());
        out.extend_from_slice(&2u64.to_le_bytes()); // 2 kvs

        let put = |out: &mut Vec<u8>, s: &str| {
            out.extend_from_slice(&(s.len() as u64).to_le_bytes());
            out.extend_from_slice(s.as_bytes());
        };
        put(&mut out, "tokenizer.ggml.tokens");
        out.extend_from_slice(&9u32.to_le_bytes()); // array
        out.extend_from_slice(&8u32.to_le_bytes()); // of string
        out.extend_from_slice(&1000u64.to_le_bytes());
        for i in 0..1000u32 {
            put(&mut out, &format!("tok{i}"));
        }
        put(&mut out, "general.name");
        out.extend_from_slice(&8u32.to_le_bytes());
        put(&mut out, "After The Array");

        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("vocab.gguf");
        std::fs::write(&path, &out).unwrap();
        let meta = read_meta(&path).unwrap();
        assert_eq!(meta.general_name.as_deref(), Some("After The Array"));
    }

    #[test]
    fn unknown_value_type_returns_partial_metadata() {
        let mut out: Vec<u8> = Vec::new();
        out.extend_from_slice(b"GGUF");
        out.extend_from_slice(&3u32.to_le_bytes());
        out.extend_from_slice(&0u64.to_le_bytes());
        out.extend_from_slice(&2u64.to_le_bytes());
        let put = |out: &mut Vec<u8>, s: &str| {
            out.extend_from_slice(&(s.len() as u64).to_le_bytes());
            out.extend_from_slice(s.as_bytes());
        };
        put(&mut out, "general.name");
        out.extend_from_slice(&8u32.to_le_bytes());
        put(&mut out, "Known Part");
        put(&mut out, "future.key");
        out.extend_from_slice(&99u32.to_le_bytes()); // unknown type
        out.extend_from_slice(&[0xAA; 16]);

        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("future.gguf");
        std::fs::write(&path, &out).unwrap();
        let meta = read_meta(&path).unwrap();
        assert_eq!(meta.general_name.as_deref(), Some("Known Part"));
    }
}