//! Just enough of the GGUF container to read a model's own name out of it.
//!
//! A `.gguf` file opens with a key/value table, and `general.name` in it is
//! what the people who published the model called it -- "Qwen3-0.6B",
//! "Qwen3.8-27B GSQ-RCO". Everything else this server knows a model by is
//! filesystem trivia: `owner/repo/file.gguf` is where it was downloaded from,
//! which is an address rather than a name, and on a phone it is a line of path
//! where a word would do.
//!
//! **Read as far as the answer and no further.** The same table holds the
//! tokenizer, which for a modern model is a 150,000-entry string array and
//! most of several megabytes; `general.*` is written first by every converter
//! in practice, so stopping at the name costs a few kilobytes instead. That is
//! what makes this affordable to run over every model in a directory, and what
//! lets the remote case work from a bounded prefix of the file rather than the
//! whole of it.
//!
//! Anything unreadable is [`None`] rather than an error, at every level. A
//! model with no name, a truncated prefix, a container version this does not
//! know and a file that is not GGUF at all are one answer here -- "this file
//! does not tell us" -- and the caller has a file name to fall back on. There
//! is nothing a reader could do with the distinction.

use std::io::Read;

/// How many bytes of a model file are worth fetching to look for its name.
///
/// Only the remote path needs a number: a local read stops when it finds the
/// key, but a file on another machine has to be asked for a fixed amount
/// before anything can be parsed. Measured 2026-09-19 against the two models
/// on this machine, `general.name` ends at byte **130** and **94** -- every
/// converter writes `general.*` before the tokenizer arrays that make up the
/// rest of the table. 8 KiB is two orders of magnitude of slack for that and
/// still makes listing a directory of models one round trip's worth of bytes
/// rather than a download, which is what decides the number: this is paid per
/// model every time a spawn screen opens.
pub const PREFIX_BYTES: u64 = 8 * 1024;

/// The longest string this will allocate for, so a corrupt length field
/// cannot ask for a gigabyte. Longer than any key or `general.*` value.
const MAX_STRING: u64 = 64 * 1024;

/// What `general.name` says, or `None` for every way of not finding out.
///
/// `read` is consumed only as far as the key: pass a file to read a local
/// model, or a cursor over a prefix to read one whose bytes came from
/// somewhere else.
pub fn name(read: &mut impl Read) -> Option<String> {
    match find(read, |key| key == "general.name") {
        Some((STRING, read)) => string(read),
        _ => None,
    }
}

/// Whether this model carries a multi-token-prediction head.
///
/// Worth asking because `llama-server` **exits** when told to use one that is
/// not there -- `--spec-type draft-mtp` on a plain model is "context type MTP
/// requested but model doesn't contain MTP layers" and then a server that
/// never comes up. So the flag can only be passed once this has said yes, and
/// a `false` here is the same answer as an unreadable file: don't ask for it.
///
/// Matched on the key's tail rather than its whole name, because the key is
/// prefixed with the architecture (`qwen35.nextn_predict_layers`) and the
/// architecture is whatever the next model is. The value is not read: a model
/// that declares the key at all is one whose tensors carry the head, and the
/// two disagreeing is a broken file rather than a state to handle.
pub fn has_mtp_head(read: &mut impl Read) -> bool {
    find(read, |key| key.ends_with(".nextn_predict_layers")).is_some()
}

/// Steps through the metadata table to the first key `wanted` accepts,
/// returning its value's type tag and the reader positioned at the value.
fn find<R: Read>(read: &mut R, wanted: impl Fn(&str) -> bool) -> Option<(u32, &mut R)> {
    let mut magic = [0u8; 4];
    read.read_exact(&mut magic).ok()?;
    if &magic != b"GGUF" {
        return None;
    }
    let _version = u32s(read)?;
    let _tensors = u64s(read)?;
    let count = u64s(read)?;
    for _ in 0..count {
        let found = string(read)?;
        let kind = u32s(read)?;
        if wanted(&found) {
            return Some((kind, read));
        }
        skip_value(kind, read)?;
    }
    None
}

// The value type tags, in the container's own numbering. Only the two this
// has to act on are named; the rest are widths, and `scalar_width` is where
// the numbering is written down once.
const STRING: u32 = 8;
const ARRAY: u32 = 9;

/// How many bytes a scalar of this type occupies, or `None` for a type that
/// is not a scalar -- which includes a tag this build does not know, since a
/// value of unknown length cannot be stepped over.
fn scalar_width(kind: u32) -> Option<u64> {
    match kind {
        // u8, i8, bool
        0 | 1 | 7 => Some(1),
        // u16, i16
        2 | 3 => Some(2),
        // u32, i32, f32
        4..=6 => Some(4),
        // u64, i64, f64
        10..=12 => Some(8),
        _ => None,
    }
}

/// Steps over one value of `kind` without keeping it.
///
/// Recursive only in the sense that an array's elements are values; GGUF
/// arrays do not nest, so the recursion is one level deep by construction.
fn skip_value(kind: u32, read: &mut impl Read) -> Option<()> {
    match kind {
        STRING => {
            let len = u64s(read)?;
            skip(len, read)
        }
        ARRAY => {
            let element = u32s(read)?;
            let count = u64s(read)?;
            match scalar_width(element) {
                // The whole array at once: this is the tokenizer's scores and
                // token types, and stepping over them one at a time is a
                // syscall per token.
                Some(width) => skip(count.checked_mul(width)?, read),
                None if element == STRING => {
                    for _ in 0..count {
                        let len = u64s(read)?;
                        skip(len, read)?;
                    }
                    Some(())
                }
                // An array of arrays, or of something this build has no width
                // for: the rest of the table can no longer be located.
                None => None,
            }
        }
        _ => skip(scalar_width(kind)?, read),
    }
}

/// Discards `count` bytes, failing if the input ends first.
///
/// Chunked against a bounded buffer rather than read into a `Vec` of the
/// stated size: the sizes here come out of the file, and the file may be a
/// truncated prefix or not a GGUF at all.
fn skip(count: u64, read: &mut impl Read) -> Option<()> {
    let mut scratch = [0u8; 8192];
    let mut left = count;
    while left > 0 {
        let want = left.min(scratch.len() as u64) as usize;
        read.read_exact(&mut scratch[..want]).ok()?;
        left -= want as u64;
    }
    Some(())
}

fn string(read: &mut impl Read) -> Option<String> {
    let len = u64s(read)?;
    if len > MAX_STRING {
        return None;
    }
    let mut bytes = vec![0u8; len as usize];
    read.read_exact(&mut bytes).ok()?;
    String::from_utf8(bytes).ok()
}

fn u32s(read: &mut impl Read) -> Option<u32> {
    let mut bytes = [0u8; 4];
    read.read_exact(&mut bytes).ok()?;
    Some(u32::from_le_bytes(bytes))
}

fn u64s(read: &mut impl Read) -> Option<u64> {
    let mut bytes = [0u8; 8];
    read.read_exact(&mut bytes).ok()?;
    Some(u64::from_le_bytes(bytes))
}

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

    /// Builds a GGUF header holding exactly these keys, so the parser is
    /// tested against the layout rather than against a fixture nobody here
    /// can regenerate.
    fn header(entries: &[(&str, Value)]) -> Vec<u8> {
        let mut out = Vec::from(*b"GGUF");
        out.extend(3u32.to_le_bytes());
        out.extend(0u64.to_le_bytes());
        out.extend((entries.len() as u64).to_le_bytes());
        for (key, value) in entries {
            put_string(&mut out, key);
            value.write(&mut out);
        }
        out
    }

    enum Value {
        Str(&'static str),
        U32(u32),
        Strings(Vec<&'static str>),
        Floats(Vec<f32>),
    }

    impl Value {
        fn write(&self, out: &mut Vec<u8>) {
            match self {
                Self::Str(text) => {
                    out.extend(STRING.to_le_bytes());
                    put_string(out, text);
                }
                Self::U32(number) => {
                    out.extend(4u32.to_le_bytes());
                    out.extend(number.to_le_bytes());
                }
                Self::Strings(items) => {
                    out.extend(ARRAY.to_le_bytes());
                    out.extend(STRING.to_le_bytes());
                    out.extend((items.len() as u64).to_le_bytes());
                    for item in items {
                        put_string(out, item);
                    }
                }
                Self::Floats(items) => {
                    out.extend(ARRAY.to_le_bytes());
                    out.extend(6u32.to_le_bytes());
                    out.extend((items.len() as u64).to_le_bytes());
                    for item in items {
                        out.extend(item.to_le_bytes());
                    }
                }
            }
        }
    }

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

    #[test]
    fn the_name_is_read_past_every_other_kind_of_value() {
        let bytes = header(&[
            ("general.architecture", Value::Str("qwen3")),
            ("general.file_type", Value::U32(7)),
            ("qwen3.attention.head_count", Value::U32(16)),
            ("tokenizer.ggml.scores", Value::Floats(vec![0.5; 64])),
            ("tokenizer.ggml.tokens", Value::Strings(vec!["a", "b", "c"])),
            ("general.name", Value::Str("Qwen3-0.6B")),
        ]);
        assert_eq!(
            name(&mut bytes.as_slice()),
            Some("Qwen3-0.6B".to_string()),
            "every value before the name has to be steppable over",
        );
    }

    #[test]
    /// The remote case: a prefix is all there is, and running off the end of
    /// it is "we don't know" rather than a failure worth reporting. The
    /// caller has the file name.
    fn a_truncated_file_has_no_name_rather_than_failing() {
        let bytes = header(&[
            ("tokenizer.ggml.tokens", Value::Strings(vec!["a", "b", "c"])),
            ("general.name", Value::Str("Qwen3-0.6B")),
        ]);
        for cut in [4, 12, 24, bytes.len() - 4] {
            assert_eq!(name(&mut &bytes[..cut]), None, "cut at {cut}");
        }
    }

    #[test]
    fn a_file_that_is_not_gguf_has_no_name() {
        assert_eq!(name(&mut b"not a model at all".as_slice()), None);
        assert_eq!(name(&mut b"".as_slice()), None);
    }

    #[test]
    /// A name that is not a string is not a name. The alternative is
    /// rendering a number as one, which reads as a model called "7".
    fn a_name_of_the_wrong_type_is_not_read() {
        let bytes = header(&[("general.name", Value::U32(7))]);
        assert_eq!(name(&mut bytes.as_slice()), None);
    }

    #[test]
    /// The head is found by the tail of the key, because the whole key is
    /// prefixed with whatever architecture the model is.
    fn an_mtp_head_is_found_whatever_the_architecture_is_called() {
        let with = header(&[
            ("general.architecture", Value::Str("qwen35")),
            ("qwen35.block_count", Value::U32(64)),
            ("qwen35.nextn_predict_layers", Value::U32(1)),
        ]);
        assert!(has_mtp_head(&mut with.as_slice()));

        let without = header(&[
            ("general.architecture", Value::Str("qwen3")),
            ("qwen3.block_count", Value::U32(28)),
        ]);
        assert!(!has_mtp_head(&mut without.as_slice()));
    }

    #[test]
    /// A prefix that stops short says no, and that is the direction it has to
    /// fail in: `--spec-type draft-mtp` on a model with no head is a server
    /// that exits, so "we could not tell" and "it has none" both mean don't
    /// ask for it.
    fn a_truncated_file_reports_no_mtp_head() {
        let bytes = header(&[("qwen35.nextn_predict_layers", Value::U32(1))]);
        assert!(!has_mtp_head(&mut &bytes[..12]));
    }

    #[test]
    /// The real thing, when this machine happens to have one. Skipped rather
    /// than failed where it does not: the models directory is not part of the
    /// checkout, and a test that needs gigabytes to run is one nobody runs.
    fn a_real_model_on_this_machine_reads_back_its_name() {
        let Some(home) = std::env::var_os("HOME") else {
            return;
        };
        let dir = std::path::Path::new(&home).join(".local/share/ai-app/models");
        let mut found = Vec::new();
        collect_gguf(&dir, &mut found);
        for path in found {
            let mut file = std::fs::File::open(&path).expect("open");
            let read = name(&mut file);
            assert!(
                read.is_some_and(|name| !name.trim().is_empty()),
                "{} has a name in it and this did not read one",
                path.display(),
            );
        }
    }

    fn collect_gguf(dir: &std::path::Path, found: &mut Vec<std::path::PathBuf>) {
        let Ok(entries) = std::fs::read_dir(dir) else {
            return;
        };
        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_dir() {
                collect_gguf(&path, found);
            } else if path.extension().is_some_and(|e| e == "gguf") {
                found.push(path);
            }
        }
    }
}