use memmap2::{Mmap, MmapOptions};
use std::collections::HashMap;
use std::fs::File;
use std::path::{Path, PathBuf};

const MAGIC: u32 = 0x4655_4747;
const MAX_COUNT: u64 = 1_000_000;
const MAX_DIMS: usize = 8;
const MAX_DEPTH: u8 = 8;

pub(super) const F32: u32 = 0;
pub(super) const F16: u32 = 1;
pub(super) const Q4_0: u32 = 2;
pub(super) const Q8_0: u32 = 8;
pub(super) const Q2_K: u32 = 10;
pub(super) const Q4_K: u32 = 12;
pub(super) const Q5_K: u32 = 13;
pub(super) const Q6_K: u32 = 14;
pub(super) const IQ2_XXS: u32 = 16;
pub(super) const I32: u32 = 26;

#[derive(Clone, Debug)]
pub(super) enum Value {
    U32(u32),
    I32(i32),
    U64(u64),
    I64(i64),
    F32(f32),
    F64(f64),
    Bool(bool),
    Bytes(Vec<u8>),
    U32s(Vec<u32>),
    F32s(Vec<f32>),
    Strings(Vec<Vec<u8>>),
    Other,
}

#[derive(Clone, Debug)]
pub(super) struct Tensor {
    pub(super) kind: u32,
    pub(super) dims: Vec<u64>,
    pub(super) offset: u64,
    pub(super) bytes: u64,
}

pub(super) struct Gguf {
    path: PathBuf,
    map: Mmap,
    data_offset: u64,
    max_tensor_bytes: u64,
    pub(super) metadata: HashMap<String, Value>,
    pub(super) tensors: HashMap<String, Tensor>,
}

impl Gguf {
    pub(super) fn open(path: &Path) -> Result<Self, String> {
        let file =
            File::open(path).map_err(|error| format!("Cannot open {}: {error}", path.display()))?;
        if file.metadata().map_err(|error| error.to_string())?.len() < 32 {
            return Err(format!("{} is too small to be GGUF", path.display()));
        }
        // SAFETY: managed model artifacts are opened read-only and are never mutated
        // while a Model owns this mapping.
        let map = unsafe { MmapOptions::new().map(&file) }
            .map_err(|error| format!("Cannot map {}: {error}", path.display()))?;
        let mut cursor = Cursor::new(&map);
        if cursor.u32()? != MAGIC {
            return Err(format!("{} is not a GGUF file", path.display()));
        }
        if cursor.u32()? != 3 {
            return Err(format!("{} is not GGUF v3", path.display()));
        }
        let tensor_count = cursor.u64()?;
        let metadata_count = cursor.u64()?;
        if tensor_count > MAX_COUNT || metadata_count > MAX_COUNT {
            return Err("GGUF directory count is too large".into());
        }

        let mut metadata = HashMap::with_capacity(metadata_count as usize);
        let mut alignment = 32_u64;
        for _ in 0..metadata_count {
            let key = String::from_utf8(cursor.string()?.to_vec())
                .map_err(|_| "GGUF metadata key is not UTF-8".to_owned())?;
            let kind = cursor.u32()?;
            let value = cursor.value(kind, &key, 0)?;
            if key == "general.alignment"
                && let Value::U32(value) = value
            {
                alignment = u64::from(value);
            }
            if metadata.insert(key.clone(), value).is_some() {
                return Err(format!("duplicate GGUF metadata key: {key}"));
            }
        }
        if alignment == 0 {
            return Err("GGUF alignment is zero".into());
        }

        let mut tensors = HashMap::with_capacity(tensor_count as usize);
        for _ in 0..tensor_count {
            let name = String::from_utf8(cursor.string()?.to_vec())
                .map_err(|_| "GGUF tensor name is not UTF-8".to_owned())?;
            let ndim = cursor.u32()? as usize;
            if ndim == 0 || ndim > MAX_DIMS {
                return Err(format!(
                    "tensor {name} has unsupported dimension count {ndim}"
                ));
            }
            let mut dims = Vec::with_capacity(ndim);
            let mut elements = 1_u64;
            for _ in 0..ndim {
                let dim = cursor.u64()?;
                elements = elements
                    .checked_mul(dim)
                    .ok_or_else(|| format!("tensor {name} element count overflows"))?;
                dims.push(dim);
            }
            let kind = cursor.u32()?;
            let relative_offset = cursor.u64()?;
            let (block_elements, block_bytes) = tensor_type(kind)
                .ok_or_else(|| format!("tensor {name} uses unsupported GGUF type {kind}"))?;
            let blocks = elements
                .checked_add(block_elements - 1)
                .ok_or_else(|| format!("tensor {name} size overflows"))?
                / block_elements;
            let bytes = blocks
                .checked_mul(block_bytes)
                .ok_or_else(|| format!("tensor {name} size overflows"))?;
            if tensors
                .insert(
                    name.clone(),
                    Tensor {
                        kind,
                        dims,
                        offset: relative_offset,
                        bytes,
                    },
                )
                .is_some()
            {
                return Err(format!("duplicate GGUF tensor: {name}"));
            }
        }

        let data_start = align(cursor.position() as u64, alignment)?;
        let mut max_tensor_bytes = 0;
        for (name, tensor) in &mut tensors {
            tensor.offset = data_start
                .checked_add(tensor.offset)
                .ok_or_else(|| format!("tensor {name} offset overflows"))?;
            let end = tensor
                .offset
                .checked_add(tensor.bytes)
                .ok_or_else(|| format!("tensor {name} end overflows"))?;
            if end > map.len() as u64 {
                return Err(format!("tensor {name} points outside the GGUF file"));
            }
            max_tensor_bytes = max_tensor_bytes.max(tensor.bytes);
        }

        Ok(Self {
            path: path.to_owned(),
            map,
            data_offset: data_start,
            max_tensor_bytes,
            metadata,
            tensors,
        })
    }

    pub(super) fn path(&self) -> &Path {
        &self.path
    }

    pub(super) fn len(&self) -> u64 {
        self.map.len() as u64
    }

    pub(super) fn map_ptr(&self) -> *const u8 {
        self.map.as_ptr()
    }

    pub(super) fn data_offset(&self) -> u64 {
        self.data_offset
    }

    pub(super) fn max_tensor_bytes(&self) -> u64 {
        self.max_tensor_bytes
    }

    pub(super) fn tensor(&self, name: &str) -> Result<&Tensor, String> {
        self.tensors
            .get(name)
            .ok_or_else(|| format!("required tensor is missing: {name}"))
    }

    pub(super) fn tensor_data(&self, name: &str) -> Result<&[u8], String> {
        let tensor = self.tensor(name)?;
        let start = tensor.offset as usize;
        let end = start + tensor.bytes as usize;
        Ok(&self.map[start..end])
    }

    pub(super) fn u32(&self, key: &str) -> Result<u32, String> {
        match self.metadata.get(key) {
            Some(Value::U32(value)) => Ok(*value),
            _ => Err(format!("required uint32 metadata key is missing: {key}")),
        }
    }

    pub(super) fn u64(&self, key: &str) -> Result<u64, String> {
        match self.metadata.get(key) {
            Some(Value::U64(value)) => Ok(*value),
            Some(Value::U32(value)) => Ok(u64::from(*value)),
            _ => Err(format!("required integer metadata key is missing: {key}")),
        }
    }

    pub(super) fn f32(&self, key: &str) -> Result<f32, String> {
        match self.metadata.get(key) {
            Some(Value::F32(value)) => Ok(*value),
            Some(Value::F64(value)) => Ok(*value as f32),
            Some(Value::U32(value)) => Ok(*value as f32),
            Some(Value::I32(value)) => Ok(*value as f32),
            _ => Err(format!("required numeric metadata key is missing: {key}")),
        }
    }

    pub(super) fn token_id(&self, key: &str) -> Result<i32, String> {
        let value = match self.metadata.get(key) {
            Some(Value::U32(value)) => i64::from(*value),
            Some(Value::I32(value)) => i64::from(*value),
            Some(Value::U64(value)) => i64::try_from(*value).unwrap_or(-1),
            Some(Value::I64(value)) => *value,
            _ => -1,
        };
        i32::try_from(value)
            .ok()
            .filter(|value| *value >= 0)
            .ok_or_else(|| format!("required tokenizer token ID is missing: {key}"))
    }

    pub(super) fn boolean(&self, key: &str) -> Result<bool, String> {
        match self.metadata.get(key) {
            Some(Value::Bool(value)) => Ok(*value),
            _ => Err(format!("required boolean metadata key is missing: {key}")),
        }
    }

    pub(super) fn bytes(&self, key: &str) -> Result<&[u8], String> {
        match self.metadata.get(key) {
            Some(Value::Bytes(value)) => Ok(value),
            _ => Err(format!("required string metadata key is missing: {key}")),
        }
    }

    pub(super) fn u32s(&self, key: &str) -> Result<&[u32], String> {
        match self.metadata.get(key) {
            Some(Value::U32s(value)) => Ok(value),
            _ => Err(format!(
                "required integer-array metadata key is missing: {key}"
            )),
        }
    }

    pub(super) fn f32s(&self, key: &str) -> Result<&[f32], String> {
        match self.metadata.get(key) {
            Some(Value::F32s(value)) => Ok(value),
            _ => Err(format!(
                "required float-array metadata key is missing: {key}"
            )),
        }
    }

    pub(super) fn strings(&self, key: &str) -> Result<&[Vec<u8>], String> {
        match self.metadata.get(key) {
            Some(Value::Strings(value)) => Ok(value),
            _ => Err(format!(
                "required string-array metadata key is missing: {key}"
            )),
        }
    }
}

fn tensor_type(kind: u32) -> Option<(u64, u64)> {
    Some(match kind {
        0 => (1, 4),
        1 => (1, 2),
        2 => (32, 18),
        3 => (32, 20),
        6 => (32, 22),
        7 => (32, 24),
        8 => (32, 34),
        9 => (32, 40),
        10 => (256, 84),
        11 => (256, 110),
        12 => (256, 144),
        13 => (256, 176),
        14 => (256, 210),
        15 => (256, 292),
        16 => (256, 66),
        17 => (256, 74),
        18 => (256, 98),
        19 => (256, 110),
        20 => (256, 50),
        21 => (256, 110),
        22 => (256, 82),
        23 => (256, 136),
        24 => (1, 1),
        25 => (1, 2),
        26 => (1, 4),
        27 => (1, 8),
        28 => (1, 8),
        29 => (256, 56),
        30 => (1, 2),
        _ => return None,
    })
}

fn align(value: u64, alignment: u64) -> Result<u64, String> {
    let remainder = value % alignment;
    if remainder == 0 {
        Ok(value)
    } else {
        value
            .checked_add(alignment - remainder)
            .ok_or_else(|| "GGUF alignment overflows".into())
    }
}

struct Cursor<'a> {
    bytes: &'a [u8],
    position: usize,
}

impl<'a> Cursor<'a> {
    fn new(bytes: &'a [u8]) -> Self {
        Self { bytes, position: 0 }
    }

    fn position(&self) -> usize {
        self.position
    }

    fn take(&mut self, count: u64) -> Result<&'a [u8], String> {
        let count = usize::try_from(count).map_err(|_| "GGUF value is too large")?;
        let end = self
            .position
            .checked_add(count)
            .filter(|end| *end <= self.bytes.len())
            .ok_or_else(|| format!("truncated GGUF at byte {}", self.position))?;
        let value = &self.bytes[self.position..end];
        self.position = end;
        Ok(value)
    }

    fn u8(&mut self) -> Result<u8, String> {
        Ok(self.take(1)?[0])
    }

    fn u16(&mut self) -> Result<u16, String> {
        Ok(u16::from_le_bytes(self.take(2)?.try_into().unwrap()))
    }

    fn u32(&mut self) -> Result<u32, String> {
        Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap()))
    }

    fn i32(&mut self) -> Result<i32, String> {
        Ok(i32::from_le_bytes(self.take(4)?.try_into().unwrap()))
    }

    fn u64(&mut self) -> Result<u64, String> {
        Ok(u64::from_le_bytes(self.take(8)?.try_into().unwrap()))
    }

    fn i64(&mut self) -> Result<i64, String> {
        Ok(i64::from_le_bytes(self.take(8)?.try_into().unwrap()))
    }

    fn f32(&mut self) -> Result<f32, String> {
        Ok(f32::from_le_bytes(self.take(4)?.try_into().unwrap()))
    }

    fn f64(&mut self) -> Result<f64, String> {
        Ok(f64::from_le_bytes(self.take(8)?.try_into().unwrap()))
    }

    fn string(&mut self) -> Result<&'a [u8], String> {
        let len = self.u64()?;
        self.take(len)
    }

    fn value(&mut self, kind: u32, key: &str, depth: u8) -> Result<Value, String> {
        if depth > MAX_DEPTH {
            return Err("GGUF metadata arrays are nested too deeply".into());
        }
        Ok(match kind {
            0 => {
                self.u8()?;
                Value::Other
            }
            1 => {
                self.u8()?;
                Value::Other
            }
            2 | 3 => {
                self.u16()?;
                Value::Other
            }
            4 => Value::U32(self.u32()?),
            5 => Value::I32(self.i32()?),
            6 => Value::F32(self.f32()?),
            7 => Value::Bool(self.u8()? != 0),
            8 => Value::Bytes(self.string()?.to_vec()),
            9 => self.array(key, depth + 1)?,
            10 => Value::U64(self.u64()?),
            11 => Value::I64(self.i64()?),
            12 => Value::F64(self.f64()?),
            _ => return Err(format!("unknown GGUF metadata type {kind}")),
        })
    }

    fn array(&mut self, key: &str, depth: u8) -> Result<Value, String> {
        let item = self.u32()?;
        let len = self.u64()?;
        if len > MAX_COUNT {
            return Err(format!("GGUF metadata array is too large: {key}"));
        }
        let keep = matches!(
            key,
            "tokenizer.ggml.tokens"
                | "tokenizer.ggml.merges"
                | "deepseek4.attention.compress_ratios"
                | "deepseek4.swiglu_clamp_exp"
                | "deepseek4.dspark.target_layer_ids"
                | "deepseek4.dspark_target_layer_ids"
                | "dspark.target_layer_ids"
        );
        if !keep {
            for _ in 0..len {
                self.value(item, key, depth)?;
            }
            return Ok(Value::Other);
        }
        match item {
            4 | 5 => {
                let mut values = Vec::with_capacity(len as usize);
                for _ in 0..len {
                    let value = if item == 4 {
                        self.u32()?
                    } else {
                        self.i32()?
                            .try_into()
                            .map_err(|_| format!("negative value in {key}"))?
                    };
                    values.push(value);
                }
                Ok(Value::U32s(values))
            }
            6 | 12 => {
                let mut values = Vec::with_capacity(len as usize);
                for _ in 0..len {
                    values.push(if item == 6 {
                        self.f32()?
                    } else {
                        self.f64()? as f32
                    });
                }
                Ok(Value::F32s(values))
            }
            8 => {
                let mut values = Vec::with_capacity(len as usize);
                for _ in 0..len {
                    values.push(self.string()?.to_vec());
                }
                Ok(Value::Strings(values))
            }
            _ => Err(format!("unsupported array type {item} for {key}")),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use std::time::{SystemTime, UNIX_EPOCH};

    #[test]
    fn maps_metadata_and_tensor_payload_without_copying_weights() {
        let path = std::env::temp_dir().join(format!(
            "ds4-server-gguf-{}",
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        ));
        let mut bytes = Vec::new();
        bytes.extend(MAGIC.to_le_bytes());
        bytes.extend(3_u32.to_le_bytes());
        bytes.extend(1_u64.to_le_bytes());
        bytes.extend(1_u64.to_le_bytes());
        push_string(&mut bytes, b"general.architecture");
        bytes.extend(8_u32.to_le_bytes());
        push_string(&mut bytes, b"deepseek4");
        push_string(&mut bytes, b"weight");
        bytes.extend(1_u32.to_le_bytes());
        bytes.extend(1_u64.to_le_bytes());
        bytes.extend(F32.to_le_bytes());
        bytes.extend(0_u64.to_le_bytes());
        bytes.resize(bytes.len().div_ceil(32) * 32, 0);
        bytes.extend(1_f32.to_le_bytes());
        fs::write(&path, bytes).unwrap();

        let model = Gguf::open(&path).unwrap();
        assert_eq!(model.bytes("general.architecture").unwrap(), b"deepseek4");
        assert_eq!(model.tensor("weight").unwrap().dims, [1]);
        assert_eq!(model.tensor_data("weight").unwrap(), 1_f32.to_le_bytes());
        fs::remove_file(path).unwrap();
    }

    fn push_string(bytes: &mut Vec<u8>, value: &[u8]) {
        bytes.extend((value.len() as u64).to_le_bytes());
        bytes.extend(value);
    }
}