use memmap2::{Advice, Mmap, MmapOptions};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::sync::Arc;

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;
pub(super) const BF16: u32 = 30;
pub(super) const MXFP4: u32 = 39;

#[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: Arc<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}"))?;
            if kind == MXFP4 && !dims[0].is_multiple_of(block_elements) {
                return Err(format!(
                    "tensor {name} MXFP4 row length {} is not aligned to {block_elements} values",
                    dims[0]
                ));
            }
            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: Arc::new(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 checkpoint_identity(&self) -> [u8; 32] {
        let mut hash = Sha256::new();
        hash.update(b"DS4Server GGUF checkpoint identity v1");
        hash.update(
            self.path
                .canonicalize()
                .unwrap_or_else(|_| self.path.clone())
                .to_string_lossy()
                .as_bytes(),
        );
        hash.update(self.len().to_le_bytes());
        hash.update(self.data_offset.to_le_bytes());
        let mut tensors = self.tensors.iter().collect::<Vec<_>>();
        tensors.sort_by_key(|(name, _)| *name);
        for (name, tensor) in tensors {
            hash.update(name.as_bytes());
            hash.update(tensor.kind.to_le_bytes());
            hash.update(tensor.offset.to_le_bytes());
            hash.update(tensor.bytes.to_le_bytes());
            for dimension in &tensor.dims {
                hash.update(dimension.to_le_bytes());
            }
        }
        hash.finalize().into()
    }

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

    pub(super) fn shared_map(&self) -> Arc<Mmap> {
        Arc::clone(&self.map)
    }

    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 warm(&self) -> Result<(), String> {
        let start = self.data_offset as usize;
        if start >= self.map.len() {
            return Ok(());
        }
        self.map
            .advise_range(Advice::WillNeed, start, self.map.len() - start)
            .map_err(|error| format!("Cannot warm {}: {error}", self.path.display()))?;
        let mut checksum = 0_u64;
        for offset in (start..self.map.len()).step_by(16 * 1024) {
            checksum = checksum.wrapping_add(u64::from(self.map[offset]));
        }
        checksum = checksum.wrapping_add(u64::from(self.map[self.map.len() - 1]));
        std::hint::black_box(checksum);
        Ok(())
    }

    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),
        39 => (32, 17),
        _ => 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"
                | "glm5-next.layer_types"
        );
        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::sync::atomic::{AtomicU64, Ordering};
    use std::time::{SystemTime, UNIX_EPOCH};

    static FIXTURE_ID: AtomicU64 = AtomicU64::new(0);

    #[test]
    fn maps_metadata_and_tensor_payload_without_copying_weights() {
        let path = write_gguf(&[1], F32, 4);

        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(), [0; 4]);
        model.warm().unwrap();
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn accepts_exact_mxfp4_block_layout() {
        let path = write_gguf(&[32, 2], MXFP4, 34);
        let model = Gguf::open(&path).unwrap();
        let tensor = model.tensor("weight").unwrap();
        assert_eq!(tensor.bytes, 34);
        assert_eq!(model.tensor_data("weight").unwrap().len(), 34);
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn rejects_mxfp4_rows_that_are_not_block_aligned() {
        let path = write_gguf(&[33, 2], MXFP4, 34);
        let error = Gguf::open(&path).err().unwrap();
        assert!(error.contains("MXFP4 row length 33 is not aligned to 32 values"));
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn rejects_truncated_mxfp4_payload() {
        let path = write_gguf(&[32, 2], MXFP4, 33);
        let error = Gguf::open(&path).err().unwrap();
        assert!(error.contains("points outside the GGUF file"));
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn rejects_mxfp4_dimension_overflow() {
        let path = write_gguf(&[32, u64::MAX], MXFP4, 0);
        let error = Gguf::open(&path).err().unwrap();
        assert!(error.contains("tensor weight element count overflows"));
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn keeps_unknown_tensor_kinds_rejected() {
        let path = write_gguf(&[32], 38, 0);
        let error = Gguf::open(&path).err().unwrap();
        assert!(error.contains("unsupported GGUF type 38"));
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn mxfp4_scalar_dot_matches_independent_values() {
        const VALUES: [f32; 16] = [
            0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0,
        ];
        let mut block = [0xf6_u8; 17];
        block[0] = 127;
        let activation = [1.0; 32];
        let scale = f32::from_bits(u32::from(block[0]) << 23);
        let mut dot = 0.0;
        for (index, packed) in block[1..].iter().copied().enumerate() {
            dot += scale * VALUES[usize::from(packed & 0x0f)] * activation[index];
            dot += scale * VALUES[usize::from(packed >> 4)] * activation[index + 16];
        }
        assert_eq!(dot, -32.0);

        for exponent in [0_u8, 1, 126, 127, 128, 254] {
            let bits = if exponent == 0 {
                0x0040_0000
            } else {
                u32::from(exponent) << 23
            };
            let expected = 2.0_f32.powi(if exponent == 0 {
                -127
            } else {
                i32::from(exponent) - 127
            });
            assert_eq!(f32::from_bits(bits), expected);
        }
    }

    fn write_gguf(dims: &[u64], kind: u32, payload_bytes: usize) -> PathBuf {
        let path = std::env::temp_dir().join(format!(
            "ds4-server-gguf-{}-{}",
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_nanos(),
            FIXTURE_ID.fetch_add(1, Ordering::Relaxed),
        ));
        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((dims.len() as u32).to_le_bytes());
        for dim in dims {
            bytes.extend(dim.to_le_bytes());
        }
        bytes.extend(kind.to_le_bytes());
        bytes.extend(0_u64.to_le_bytes());
        bytes.resize(bytes.len().div_ceil(32) * 32 + payload_bytes, 0);
        fs::write(&path, bytes).unwrap();
        path
    }

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