/// .mvec binary file format — MonaVec index serialization.
///
/// Layout v4 (little-endian, 56 bytes):
///   [MAGIC:          4B ] b"MVEC"
///   [VERSION:        4B ] u32 = 4
///   [DIM:            4B ] u32
///   [METRIC:         1B ] u8  (0=Cosine, 1=Dot, 2=L2)
///   [BIT_WIDTH:      1B ] u8  (2 or 4, dominant mode)
///   [INDEX_TYPE:     1B ] u8  (0=BruteForce, 1=IvfFlat, 2=Hnsw)
///   [_PAD:           1B ] reserved
///   [COUNT:          8B ] u64 — number of indexed vectors
///   [SEED:           8B ] u64 — ChaCha20 seed for RHDH rotation
///   [N4_DIMS:        4B ] u32 — dims at 4-bit in padded space
///   [_PAD4:          4B ] reserved
///   [IVF_NLIST:      4B ] u32 — IvfFlat: number of clusters (0 for other types)
///   [IVF_NPROBE:     4B ] u32 — IvfFlat: cells probed per query
///   [HNSW_M:         2B ] u16 — HNSW: max neighbours per node
///   [HNSW_EF_CONS:   2B ] u16 — HNSW: ef_construction
///   [HNSW_EF_SEARCH: 2B ] u16 — HNSW: ef_search
///   [_PAD2:          2B ] reserved
///   [VECTORS:        u8[]  ] packed quantized data, COUNT × packed_bytes
///   [IDS:            u64[] ] vector IDs, COUNT × 8 bytes
///   [NORMS:          f32[] ] per-vector norms (v2+), COUNT × 4 bytes
///   [IVF_DATA:       ...   ] IvfFlat centroid + list data (v4, only when INDEX_TYPE=1)
///
/// v1/v2: 32-byte header, no N4_DIMS.
/// v3:    40-byte header, no index params.
/// v4:    56-byte header, full index params + IVF_DATA section for IvfFlat.

use std::io::{self, Read, Write};
use crate::{Config, IndexType, Metric, error::{MonaVecError, Result}};

pub const MAGIC:        &[u8; 4] = b"MVEC";
pub const VERSION:      u32      = 6;
pub const BYTE_SIZE_V2: usize    = 32; // v1 and v2: no N4_DIMS
pub const BYTE_SIZE_V3: usize    = 40; // v3: with N4_DIMS + pad4, no index params
pub const BYTE_SIZE_V4: usize    = 56; // v4/v5/v6: + IvfFlat/HNSW params (16 bytes, _PAD2 reused)

/// Header that fits in one fixed-size block.
#[derive(Debug, Clone)]
pub struct Header {
    pub dim:                  u32,
    pub metric:               u8,
    pub bit_width:            u8,
    pub index_type:           u8,
    pub count:                u64,
    pub seed:                 u64,
    /// Dimensions at 4-bit in padded space. For v1/v2 files, derived from bit_width.
    pub n4_dims:              u32,
    /// File format version read from disk; always writes current VERSION.
    pub version:              u32,
    // IvfFlat params (v4+; 0 for BruteForce/HNSW)
    pub ivf_nlist:            u32,
    pub ivf_nprobe:           u32,
    // HNSW params (v4+; 0 for BruteForce/IvfFlat)
    pub hnsw_m:               u16,
    pub hnsw_ef_construction: u16,
    pub hnsw_ef_search:       u16,
    /// v6+: global-scalar std params follow immediately after the 56-byte fixed header.
    /// v1–v5 files: always false.
    pub has_std:              bool,
}

impl Header {
    /// Byte size of the full header on disk, including optional std params.
    ///
    /// v1/v2 = 32B, v3 = 40B, v4/v5/v6 fixed = 56B.
    /// v6 with has_std: fixed + dim × 8 (mean f32[] + inv_std f32[]).
    pub fn byte_size(&self) -> usize {
        let fixed = if self.version <= 2 {
            BYTE_SIZE_V2
        } else if self.version == 3 {
            BYTE_SIZE_V3
        } else {
            BYTE_SIZE_V4
        };
        if self.has_std { fixed + self.dim as usize * 8 } else { fixed }
    }
}

impl Header {
    pub fn from_config(config: &Config, count: usize) -> Self {
        Self {
            dim:                  config.dim as u32,
            metric:               config.metric.as_u8(),
            bit_width:            config.bit_width,
            index_type:           config.index_type.as_u8(),
            count:                count as u64,
            seed:                 config.seed,
            n4_dims:              config.n4_dims as u32,
            version:              VERSION,
            ivf_nlist:            config.ivf_nlist  as u32,
            ivf_nprobe:           config.ivf_nprobe as u32,
            hnsw_m:               config.hnsw_m               as u16,
            hnsw_ef_construction: config.hnsw_ef_construction as u16,
            hnsw_ef_search:       config.hnsw_ef_search       as u16,
            has_std:              false,  // engine sets to true when pipeline has std_params
        }
    }

    pub fn to_config(&self) -> Result<Config> {
        let metric     = Metric::from_u8(self.metric)?;
        let index_type = IndexType::from_u8(self.index_type)?;
        let dim        = self.dim as usize;
        // For v3 files: n4_dims is stored explicitly.
        // For v1/v2 files: derive from bit_width.
        let n4_dims = if self.version >= 3 {
            self.n4_dims as usize
        } else {
            match self.bit_width {
                4 => dim.next_power_of_two(),
                _ => 0,
            }
        };
        Ok(Config {
            dim,
            metric,
            bit_width:            self.bit_width,
            n4_dims,
            seed:                 self.seed,
            index_type,
            ivf_nlist:            self.ivf_nlist  as usize,
            ivf_nprobe:           self.ivf_nprobe as usize,
            hnsw_m:               self.hnsw_m     as usize,
            hnsw_ef_construction: self.hnsw_ef_construction as usize,
            hnsw_ef_search:       self.hnsw_ef_search       as usize,
        })
    }

    pub fn write<W: Write>(&self, w: &mut W) -> io::Result<()> {
        w.write_all(MAGIC)?;
        w.write_all(&VERSION.to_le_bytes())?;
        w.write_all(&self.dim.to_le_bytes())?;
        w.write_all(&[self.metric, self.bit_width, self.index_type, 0])?;
        w.write_all(&self.count.to_le_bytes())?;
        w.write_all(&self.seed.to_le_bytes())?;
        // v3+: N4_DIMS + 4 bytes padding
        w.write_all(&self.n4_dims.to_le_bytes())?;
        w.write_all(&0u32.to_le_bytes())?;
        // v4: IvfFlat / HNSW params (16 bytes)
        w.write_all(&self.ivf_nlist.to_le_bytes())?;
        w.write_all(&self.ivf_nprobe.to_le_bytes())?;
        w.write_all(&self.hnsw_m.to_le_bytes())?;
        w.write_all(&self.hnsw_ef_construction.to_le_bytes())?;
        w.write_all(&self.hnsw_ef_search.to_le_bytes())?;
        // _PAD2: first byte is has_std flag (v6+); second byte reserved.
        w.write_all(&[self.has_std as u8, 0u8])?;
        Ok(())
    }

    pub fn read<R: Read>(r: &mut R) -> Result<Self> {
        let mut magic = [0u8; 4];
        r.read_exact(&mut magic).map_err(MonaVecError::Io)?;
        if &magic != MAGIC {
            return Err(MonaVecError::InvalidFormat("bad magic bytes".into()));
        }

        let mut version_buf = [0u8; 4];
        r.read_exact(&mut version_buf).map_err(MonaVecError::Io)?;
        let version = u32::from_le_bytes(version_buf);
        if version == 0 || version > VERSION {
            return Err(MonaVecError::UnsupportedVersion(version));
        }

        let mut dim_buf = [0u8; 4];
        r.read_exact(&mut dim_buf).map_err(MonaVecError::Io)?;
        let dim = u32::from_le_bytes(dim_buf);

        let mut flags = [0u8; 4]; // metric, bit_width, index_type, pad
        r.read_exact(&mut flags).map_err(MonaVecError::Io)?;
        let metric     = flags[0];
        let bit_width  = flags[1];
        let index_type = flags[2];

        let mut count_buf = [0u8; 8];
        r.read_exact(&mut count_buf).map_err(MonaVecError::Io)?;
        let count = u64::from_le_bytes(count_buf);

        let mut seed_buf = [0u8; 8];
        r.read_exact(&mut seed_buf).map_err(MonaVecError::Io)?;
        let seed = u64::from_le_bytes(seed_buf);

        // v3+: read N4_DIMS + pad4 (8 bytes total)
        let n4_dims = if version >= 3 {
            let mut buf = [0u8; 4];
            r.read_exact(&mut buf).map_err(MonaVecError::Io)?;
            let n4 = u32::from_le_bytes(buf);
            r.read_exact(&mut buf).map_err(MonaVecError::Io)?; // pad4, ignored
            n4
        } else {
            0 // sentinel; to_config() overrides for v1/v2
        };

        // v4+: IvfFlat / HNSW params (16 bytes).
        // Last 2 bytes (_PAD2): for v6+ first byte is has_std flag; v1–v5 always 0.
        let (ivf_nlist, ivf_nprobe, hnsw_m, hnsw_ef_construction, hnsw_ef_search, has_std_raw) =
            if version >= 4 {
                let mut b4 = [0u8; 4];
                let mut b2 = [0u8; 2];
                r.read_exact(&mut b4).map_err(MonaVecError::Io)?;
                let nlist = u32::from_le_bytes(b4);
                r.read_exact(&mut b4).map_err(MonaVecError::Io)?;
                let nprobe = u32::from_le_bytes(b4);
                r.read_exact(&mut b2).map_err(MonaVecError::Io)?;
                let m = u16::from_le_bytes(b2);
                r.read_exact(&mut b2).map_err(MonaVecError::Io)?;
                let ef_cons = u16::from_le_bytes(b2);
                r.read_exact(&mut b2).map_err(MonaVecError::Io)?;
                let ef_search = u16::from_le_bytes(b2);
                let mut pad2 = [0u8; 2];
                r.read_exact(&mut pad2).map_err(MonaVecError::Io)?;
                (nlist, nprobe, m, ef_cons, ef_search, pad2[0])
            } else {
                (0, 0, 0, 0, 0, 0)
            };
        let has_std = version >= 6 && has_std_raw != 0;

        Ok(Self {
            dim, metric, bit_width, index_type, count, seed, n4_dims, version,
            ivf_nlist,
            ivf_nprobe,
            hnsw_m,
            hnsw_ef_construction,
            hnsw_ef_search,
            has_std,
        })
    }
}

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

    #[test]
    fn header_roundtrip() {
        let config = Config::new(1536, Metric::Cosine, 4);
        let header = Header::from_config(&config, 42);
        let mut buf = Vec::new();
        header.write(&mut buf).unwrap();

        assert_eq!(buf.len(), BYTE_SIZE_V4);

        let restored = Header::read(&mut Cursor::new(&buf)).unwrap();
        assert_eq!(restored.dim,       1536);
        assert_eq!(restored.bit_width, 4);
        assert_eq!(restored.count,     42);
        assert_eq!(restored.seed,      config.seed);
        assert_eq!(restored.n4_dims,   1536usize.next_power_of_two() as u32);
    }

    #[test]
    fn mixed_header_roundtrip() {
        let config = Config::new(384, Metric::Cosine, 4).with_mixed_ratio(0.5);
        assert!(config.is_mixed());
        let header = Header::from_config(&config, 10);
        let mut buf = Vec::new();
        header.write(&mut buf).unwrap();

        let restored = Header::read(&mut Cursor::new(&buf)).unwrap();
        assert_eq!(restored.n4_dims, config.n4_dims as u32);
        let rc = restored.to_config().unwrap();
        assert_eq!(rc.n4_dims, config.n4_dims);
    }

    #[test]
    fn bad_magic_rejected() {
        let mut buf = b"XXXX".to_vec();
        buf.extend_from_slice(&[0u8; 36]);
        let result = Header::read(&mut Cursor::new(&buf));
        assert!(matches!(result, Err(MonaVecError::InvalidFormat(_))));
    }

    #[test]
    fn future_version_rejected() {
        let config = Config::new(64, Metric::Dot, 4);
        let header = Header::from_config(&config, 0);
        let mut buf = Vec::new();
        header.write(&mut buf).unwrap();
        buf[4..8].copy_from_slice(&99u32.to_le_bytes());
        let result = Header::read(&mut Cursor::new(&buf));
        assert!(matches!(result, Err(MonaVecError::UnsupportedVersion(99))));
    }

    #[test]
    fn v1_file_accepted() {
        let config = Config::new(64, Metric::Cosine, 4);
        let header = Header::from_config(&config, 0);
        let mut buf = Vec::new();
        header.write(&mut buf).unwrap();
        buf[4..8].copy_from_slice(&1u32.to_le_bytes());
        buf.truncate(BYTE_SIZE_V2);
        let result = Header::read(&mut Cursor::new(&buf));
        assert!(result.is_ok());
        assert_eq!(result.unwrap().version, 1);
    }
}
