//! GGUF file format parser.
//!
//! Parses header, metadata key-value pairs, and tensor descriptors
//! from GGUF v2/v3 files. Does NOT read tensor data — that's done
//! via memory-mapped weights (MappedWeights).

use byteorder::{LittleEndian, ReadBytesExt};
use std::collections::HashMap;
use std::fs::File;
use std::io::{BufReader, Read, Seek};
use std::path::Path;
use thiserror::Error;

const GGUF_MAGIC: u32 = 0x46554747; // "GGUF" as little-endian u32: [G=0x47, G=0x47, U=0x55, F=0x46]

#[derive(Error, Debug)]
pub enum GgufError {
    #[error("Invalid GGUF magic: expected 0x{:08X}, got 0x{got:08X}", GGUF_MAGIC)]
    BadMagic { got: u32 },
    #[error("Unsupported GGUF version: {0} (only 2 and 3 supported)")]
    UnsupportedVersion(u32),
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),
    #[error("Invalid metadata value type: {0}")]
    InvalidValueType(u32),
    #[error("Config error: {0}")]
    ConfigError(String),
}

/// GGML tensor data types.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GgmlType {
    F32 = 0,
    F16 = 1,
    Q4_0 = 2,
    Q4_1 = 3,
    Q5_0 = 6,
    Q5_1 = 7,
    Q8_0 = 8,
    Q8_1 = 9,
    Q2K = 10,
    Q3K = 11,
    Q4K = 12,
    Q5K = 13,
    Q6K = 14,
    Iq2Xxs = 16,
    Iq2Xs = 17,
    Iq3Xxs = 18,
    Iq1S = 19,
    Iq4Nl = 20,
    Iq3S = 21,
    Iq2S = 22,
    Iq4Xs = 23,
    I8 = 24,
    I16 = 25,
    I32 = 26,
    I64 = 27,
    F64 = 28,
    Iq1M = 29,
}

impl GgmlType {
    pub fn from_u32(v: u32) -> Option<Self> {
        match v {
            0 => Some(Self::F32),
            1 => Some(Self::F16),
            2 => Some(Self::Q4_0),
            3 => Some(Self::Q4_1),
            6 => Some(Self::Q5_0),
            7 => Some(Self::Q5_1),
            8 => Some(Self::Q8_0),
            9 => Some(Self::Q8_1),
            10 => Some(Self::Q2K),
            11 => Some(Self::Q3K),
            12 => Some(Self::Q4K),
            13 => Some(Self::Q5K),
            14 => Some(Self::Q6K),
            16 => Some(Self::Iq2Xxs),
            17 => Some(Self::Iq2Xs),
            18 => Some(Self::Iq3Xxs),
            19 => Some(Self::Iq1S),
            20 => Some(Self::Iq4Nl),
            21 => Some(Self::Iq3S),
            22 => Some(Self::Iq2S),
            23 => Some(Self::Iq4Xs),
            24 => Some(Self::I8),
            25 => Some(Self::I16),
            26 => Some(Self::I32),
            27 => Some(Self::I64),
            28 => Some(Self::F64),
            29 => Some(Self::Iq1M),
            _ => None,
        }
    }
}

/// Metadata value types.
#[derive(Debug, Clone)]
pub enum MetadataValue {
    U8(u8),
    I8(i8),
    U16(u16),
    I16(i16),
    U32(u32),
    I32(i32),
    F32(f32),
    Bool(bool),
    String(String),
    Array(Vec<MetadataValue>),
    U64(u64),
    I64(i64),
    F64(f64),
}

/// Information about one tensor in the GGUF file.
#[derive(Debug, Clone)]
pub struct TensorInfo {
    pub name: String,
    pub dimensions: Vec<u64>,
    pub dtype: GgmlType,
    pub offset: u64,
}

/// Parsed GGUF header (metadata + tensor info, no tensor data).
#[derive(Debug)]
pub struct GgufHeader {
    pub version: u32,
    pub tensor_count: u64,
    pub metadata: HashMap<String, MetadataValue>,
    pub tensors: Vec<TensorInfo>,
    pub data_offset: u64,
}

/// Parse a GGUF file header.
pub fn parse_header(path: &Path) -> Result<GgufHeader, GgufError> {
    let file = File::open(path)?;
    let mut reader = BufReader::new(file);

    // Magic
    let magic = reader.read_u32::<LittleEndian>()?;
    if magic != GGUF_MAGIC {
        return Err(GgufError::BadMagic { got: magic });
    }

    // Version
    let version = reader.read_u32::<LittleEndian>()?;
    if version != 2 && version != 3 {
        return Err(GgufError::UnsupportedVersion(version));
    }

    // Counts
    let tensor_count = reader.read_u64::<LittleEndian>()?;
    let metadata_count = reader.read_u64::<LittleEndian>()?;

    // Metadata
    let mut metadata = HashMap::new();
    for _ in 0..metadata_count {
        let key = read_string(&mut reader)?;
        let value = read_metadata_value(&mut reader)?;
        metadata.insert(key, value);
    }

    // Tensor info
    let mut tensors = Vec::with_capacity(tensor_count as usize);
    for _ in 0..tensor_count {
        let name = read_string(&mut reader)?;
        let n_dims = reader.read_u32::<LittleEndian>()?;
        let mut dimensions = Vec::with_capacity(n_dims as usize);
        for _ in 0..n_dims {
            dimensions.push(reader.read_u64::<LittleEndian>()?);
        }
        let dtype_raw = reader.read_u32::<LittleEndian>()?;
        let dtype =
            GgmlType::from_u32(dtype_raw).ok_or(GgufError::InvalidValueType(dtype_raw))?;
        let offset = reader.read_u64::<LittleEndian>()?;

        tensors.push(TensorInfo {
            name,
            dimensions,
            dtype,
            offset,
        });
    }

    // Data offset: align current position to 32 bytes
    let current_pos = reader.stream_position()?;
    let data_offset = current_pos.div_ceil(32) * 32;

    Ok(GgufHeader {
        version,
        tensor_count,
        metadata,
        tensors,
        data_offset,
    })
}

fn read_string<R: Read>(reader: &mut R) -> Result<String, GgufError> {
    let len = reader.read_u64::<LittleEndian>()? as usize;
    let mut buf = vec![0u8; len];
    reader.read_exact(&mut buf)?;
    Ok(String::from_utf8_lossy(&buf).to_string())
}

fn read_metadata_value<R: Read>(reader: &mut R) -> Result<MetadataValue, GgufError> {
    let vtype = reader.read_u32::<LittleEndian>()?;
    match vtype {
        0 => Ok(MetadataValue::U8(reader.read_u8()?)),
        1 => Ok(MetadataValue::I8(reader.read_i8()?)),
        2 => Ok(MetadataValue::U16(reader.read_u16::<LittleEndian>()?)),
        3 => Ok(MetadataValue::I16(reader.read_i16::<LittleEndian>()?)),
        4 => Ok(MetadataValue::U32(reader.read_u32::<LittleEndian>()?)),
        5 => Ok(MetadataValue::I32(reader.read_i32::<LittleEndian>()?)),
        6 => Ok(MetadataValue::F32(reader.read_f32::<LittleEndian>()?)),
        7 => {
            let val = reader.read_u8()?;
            Ok(MetadataValue::Bool(val != 0))
        }
        8 => {
            let s = read_string(reader)?;
            Ok(MetadataValue::String(s))
        }
        9 => {
            let elem_type = reader.read_u32::<LittleEndian>()?;
            let count = reader.read_u64::<LittleEndian>()? as usize;
            let mut arr = Vec::with_capacity(count);
            for _ in 0..count {
                // For arrays, we need to read values without the type prefix
                let val = read_array_element(reader, elem_type)?;
                arr.push(val);
            }
            Ok(MetadataValue::Array(arr))
        }
        10 => Ok(MetadataValue::U64(reader.read_u64::<LittleEndian>()?)),
        11 => Ok(MetadataValue::I64(reader.read_i64::<LittleEndian>()?)),
        12 => Ok(MetadataValue::F64(reader.read_f64::<LittleEndian>()?)),
        _ => Err(GgufError::InvalidValueType(vtype)),
    }
}

fn read_array_element<R: Read>(
    reader: &mut R,
    elem_type: u32,
) -> Result<MetadataValue, GgufError> {
    match elem_type {
        0 => Ok(MetadataValue::U8(reader.read_u8()?)),
        1 => Ok(MetadataValue::I8(reader.read_i8()?)),
        2 => Ok(MetadataValue::U16(reader.read_u16::<LittleEndian>()?)),
        3 => Ok(MetadataValue::I16(reader.read_i16::<LittleEndian>()?)),
        4 => Ok(MetadataValue::U32(reader.read_u32::<LittleEndian>()?)),
        5 => Ok(MetadataValue::I32(reader.read_i32::<LittleEndian>()?)),
        6 => Ok(MetadataValue::F32(reader.read_f32::<LittleEndian>()?)),
        7 => Ok(MetadataValue::Bool(reader.read_u8()? != 0)),
        8 => Ok(MetadataValue::String(read_string(reader)?)),
        10 => Ok(MetadataValue::U64(reader.read_u64::<LittleEndian>()?)),
        11 => Ok(MetadataValue::I64(reader.read_i64::<LittleEndian>()?)),
        12 => Ok(MetadataValue::F64(reader.read_f64::<LittleEndian>()?)),
        _ => Err(GgufError::InvalidValueType(elem_type)),
    }
}

// --- Helper functions for reading metadata ---

pub fn get_string(header: &GgufHeader, key: &str) -> Option<String> {
    match header.metadata.get(key)? {
        MetadataValue::String(s) => Some(s.clone()),
        _ => None,
    }
}

pub fn get_u32(header: &GgufHeader, key: &str) -> Option<u32> {
    match header.metadata.get(key)? {
        MetadataValue::U32(v) => Some(*v),
        _ => None,
    }
}

pub fn get_f32(header: &GgufHeader, key: &str) -> Option<f32> {
    match header.metadata.get(key)? {
        MetadataValue::F32(v) => Some(*v),
        _ => None,
    }
}

// --- Memory-mapped weights ---

use memmap2::Mmap;

/// Memory-mapped weight data from a GGUF file.
pub struct MappedWeights {
    mmap: Mmap,
    data_offset: u64,
}

impl MappedWeights {
    pub fn open(path: &Path, data_offset: u64) -> Result<Self, GgufError> {
        let file = File::open(path)?;
        let mmap = unsafe { Mmap::map(&file)? };
        Ok(Self { mmap, data_offset })
    }

    pub fn tensor_bytes(&self, offset: u64, byte_len: usize) -> Result<&[u8], GgufError> {
        let start = (self.data_offset + offset) as usize;
        let end = start + byte_len;
        if end > self.mmap.len() {
            return Err(GgufError::Io(std::io::Error::new(
                std::io::ErrorKind::UnexpectedEof,
                format!(
                    "Tensor data out of bounds: offset {start}..{end} exceeds file size {}",
                    self.mmap.len()
                ),
            )));
        }
        Ok(&self.mmap[start..end])
    }
}

// --- Test helper: build a minimal valid GGUF for testing ---

#[cfg(test)]
pub fn build_minimal_gguf() -> Vec<u8> {
    use byteorder::WriteBytesExt;
    use std::io::Write;

    let mut buf: Vec<u8> = Vec::new();

    // Magic + version 3
    buf.write_u32::<LittleEndian>(GGUF_MAGIC).unwrap();
    buf.write_u32::<LittleEndian>(3).unwrap();

    // tensor_count=1, metadata_count=2
    buf.write_u64::<LittleEndian>(1).unwrap();
    buf.write_u64::<LittleEndian>(2).unwrap();

    // Metadata: general.architecture = "llama"
    let key = "general.architecture";
    buf.write_u64::<LittleEndian>(key.len() as u64).unwrap();
    buf.write_all(key.as_bytes()).unwrap();
    buf.write_u32::<LittleEndian>(8).unwrap(); // String type
    let val = "llama";
    buf.write_u64::<LittleEndian>(val.len() as u64).unwrap();
    buf.write_all(val.as_bytes()).unwrap();

    // Metadata: llama.block_count = 2
    let key2 = "llama.block_count";
    buf.write_u64::<LittleEndian>(key2.len() as u64).unwrap();
    buf.write_all(key2.as_bytes()).unwrap();
    buf.write_u32::<LittleEndian>(4).unwrap(); // U32 type
    buf.write_u32::<LittleEndian>(2).unwrap();

    // Tensor info: "test.weight" [4, 4] F32
    let tname = "test.weight";
    buf.write_u64::<LittleEndian>(tname.len() as u64).unwrap();
    buf.write_all(tname.as_bytes()).unwrap();
    buf.write_u32::<LittleEndian>(2).unwrap(); // 2 dims
    buf.write_u64::<LittleEndian>(4).unwrap(); // dim 0
    buf.write_u64::<LittleEndian>(4).unwrap(); // dim 1
    buf.write_u32::<LittleEndian>(0).unwrap(); // F32
    buf.write_u64::<LittleEndian>(0).unwrap(); // offset

    // Align to 32 bytes
    let current = buf.len();
    let aligned = current.div_ceil(32) * 32;
    buf.resize(aligned, 0);

    // Write 16 f32 values (4x4 tensor)
    for i in 0..16 {
        buf.write_f32::<LittleEndian>(i as f32 * 0.1).unwrap();
    }

    buf
}

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

    #[test]
    fn test_parse_minimal_gguf() {
        let data = build_minimal_gguf();
        let tmp = tempfile::NamedTempFile::new().unwrap();
        std::fs::write(tmp.path(), &data).unwrap();

        let header = parse_header(tmp.path()).unwrap();
        assert_eq!(header.version, 3);
        assert_eq!(header.tensor_count, 1);
        assert_eq!(
            get_string(&header, "general.architecture"),
            Some("llama".into())
        );
        assert_eq!(get_u32(&header, "llama.block_count"), Some(2));
        assert_eq!(header.tensors.len(), 1);
        assert_eq!(header.tensors[0].name, "test.weight");
        assert_eq!(header.tensors[0].dimensions, vec![4, 4]);
        assert_eq!(header.tensors[0].dtype, GgmlType::F32);
    }

    #[test]
    fn test_bad_magic_rejected() {
        let mut data = build_minimal_gguf();
        data[0] = 0xFF; // corrupt magic
        let tmp = tempfile::NamedTempFile::new().unwrap();
        std::fs::write(tmp.path(), &data).unwrap();

        let result = parse_header(tmp.path());
        assert!(matches!(result, Err(GgufError::BadMagic { .. })));
    }

    #[test]
    fn test_unsupported_version_rejected() {
        let mut data = build_minimal_gguf();
        // Version is at bytes 4-7, set to 99
        data[4] = 99;
        data[5] = 0;
        data[6] = 0;
        data[7] = 0;
        let tmp = tempfile::NamedTempFile::new().unwrap();
        std::fs::write(tmp.path(), &data).unwrap();

        let result = parse_header(tmp.path());
        assert!(matches!(result, Err(GgufError::UnsupportedVersion(99))));
    }

    #[test]
    fn test_ggml_type_from_u32() {
        assert_eq!(GgmlType::from_u32(0), Some(GgmlType::F32));
        assert_eq!(GgmlType::from_u32(1), Some(GgmlType::F16));
        assert_eq!(GgmlType::from_u32(8), Some(GgmlType::Q8_0));
        assert_eq!(GgmlType::from_u32(12), Some(GgmlType::Q4K));
        assert_eq!(GgmlType::from_u32(255), None);
    }

    #[test]
    fn test_get_helpers_return_none_for_missing() {
        let header = GgufHeader {
            version: 3,
            tensor_count: 0,
            metadata: HashMap::new(),
            tensors: vec![],
            data_offset: 0,
        };
        assert_eq!(get_string(&header, "missing"), None);
        assert_eq!(get_u32(&header, "missing"), None);
    }

    #[test]
    fn test_mapped_weights_open_valid_file() {
        let data = build_minimal_gguf();
        let tmp = tempfile::NamedTempFile::new().unwrap();
        std::fs::write(tmp.path(), &data).unwrap();
        let header = parse_header(tmp.path()).unwrap();
        let mapped = MappedWeights::open(tmp.path(), header.data_offset).unwrap();
        assert!(!mapped.mmap.is_empty());
    }
}
