// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use crate::distance::{fvec_distance, fvec_normalize, MetricType};
use crate::index_io_util::validate_reserved_zero;
use crate::io::{PreadCursor, ReadRequest, SeekRead, SeekWrite};
use crate::ivfflat::IVFFlatIndex;
use crate::ivfpq::RowIdFilter;
use crate::kmeans;
use roaring::RoaringTreemap;
use std::io;

pub const IVFFLAT_MAGIC: u32 = 0x4956464C; // "IVFL"
pub const IVFFLAT_VERSION: u32 = 1;
pub const IVFFLAT_HEADER_SIZE: usize = 64;

const FLAG_DELTA_IDS: u32 = 1 << 0;
const REQUIRED_FLAGS: u32 = FLAG_DELTA_IDS;
const SUPPORTED_FLAGS: u32 = REQUIRED_FLAGS;

pub fn write_ivfflat_index(index: &IVFFlatIndex, out: &mut dyn SeekWrite) -> io::Result<()> {
    let d = index.d;
    let nlist = index.nlist;
    validate_index_shape(index)?;
    let d_i32 = usize_to_i32(d, "dimension")?;
    let nlist_i32 = usize_to_i32(nlist, "nlist")?;
    let total_vectors = index.ids.iter().try_fold(0i64, |sum, ids| {
        let count = usize_to_i64(ids.len(), "total vector count")?;
        sum.checked_add(count).ok_or_else(|| {
            io::Error::new(
                io::ErrorKind::InvalidInput,
                "total vector count exceeds i64 length limit",
            )
        })
    })?;

    let mut sorted_lists: Vec<(Vec<i64>, Vec<u8>, Vec<f32>)> = Vec::with_capacity(nlist);
    for list_id in 0..nlist {
        let count = index.ids[list_id].len();
        if count == 0 {
            sorted_lists.push((Vec::new(), Vec::new(), Vec::new()));
            continue;
        }

        let mut order: Vec<usize> = (0..count).collect();
        order.sort_by_key(|&idx| index.ids[list_id][idx]);

        let sorted_ids: Vec<i64> = order.iter().map(|&idx| index.ids[list_id][idx]).collect();
        let mut sorted_vectors = Vec::with_capacity(count * d);
        for idx in order {
            sorted_vectors.extend_from_slice(&index.vectors[list_id][idx * d..(idx + 1) * d]);
        }
        let (_, id_bytes) = encode_delta_varint_ids(&sorted_ids);
        sorted_lists.push((sorted_ids, id_bytes, sorted_vectors));
    }

    write_u32_le(out, IVFFLAT_MAGIC)?;
    write_u32_le(out, IVFFLAT_VERSION)?;
    write_i32_le(out, d_i32)?;
    write_i32_le(out, nlist_i32)?;
    write_u32_le(out, index.metric as u32)?;
    write_i64_le(out, total_vectors)?;
    write_u32_le(out, FLAG_DELTA_IDS)?;
    out.write_all(&[0u8; 32])?;

    write_f32_slice(out, &index.quantizer_centroids)?;

    let offset_table_size = nlist.checked_mul(16).ok_or_else(|| {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            "IVF-FLAT offset table size overflow",
        )
    })?;
    let data_start = out
        .pos()
        .checked_add(offset_table_size as u64)
        .ok_or_else(|| {
            io::Error::new(
                io::ErrorKind::InvalidInput,
                "IVF-FLAT data start offset overflow",
            )
        })?;
    let mut list_offsets = vec![0i64; nlist];
    let mut list_counts = vec![0i32; nlist];
    let mut list_id_bytes_lens = vec![0i32; nlist];
    let mut current_offset = data_start;

    for list_id in 0..nlist {
        list_offsets[list_id] = u64_to_i64(current_offset, "list offset")?;
        let count = sorted_lists[list_id].0.len();
        list_counts[list_id] = usize_to_i32(count, "list count")?;
        if count > 0 {
            let id_bytes_len = sorted_lists[list_id].1.len();
            list_id_bytes_lens[list_id] = usize_to_i32(id_bytes_len, "delta ID section")?;
            let vector_bytes = checked_list_bytes(
                count,
                d.checked_mul(4).ok_or_else(|| {
                    io::Error::new(
                        io::ErrorKind::InvalidInput,
                        "IVF-FLAT bytes per vector overflow",
                    )
                })?,
            )?;
            let list_bytes = 12usize
                .checked_add(id_bytes_len)
                .and_then(|len| len.checked_add(vector_bytes))
                .ok_or_else(|| {
                    io::Error::new(io::ErrorKind::InvalidInput, "IVF-FLAT list size overflow")
                })?;
            current_offset = current_offset
                .checked_add(list_bytes as u64)
                .ok_or_else(|| {
                    io::Error::new(io::ErrorKind::InvalidInput, "IVF-FLAT offset overflow")
                })?;
        }
    }

    for list_id in 0..nlist {
        write_i64_le(out, list_offsets[list_id])?;
        write_i32_le(out, list_counts[list_id])?;
        write_i32_le(out, list_id_bytes_lens[list_id])?;
    }

    for (sorted_ids, id_bytes, sorted_vectors) in sorted_lists {
        if sorted_ids.is_empty() {
            continue;
        }
        write_i64_le(out, sorted_ids[0])?;
        write_i32_le(out, id_bytes.len() as i32)?;
        out.write_all(&id_bytes)?;
        write_f32_slice(out, &sorted_vectors)?;
    }

    Ok(())
}

pub struct IVFFlatIndexReader<R: SeekRead> {
    reader: R,
    pub d: usize,
    pub nlist: usize,
    pub metric: MetricType,
    pub total_vectors: i64,
    pub quantizer_centroids: Vec<f32>,
    pub list_offsets: Vec<i64>,
    pub list_counts: Vec<i32>,
    pub list_id_bytes_lens: Vec<i32>,
    delta_ids: bool,
    loaded: bool,
}

impl<R: SeekRead> IVFFlatIndexReader<R> {
    pub fn open(mut reader: R) -> io::Result<Self> {
        let mut cursor = PreadCursor::new(&mut reader, 0);

        let magic = read_u32_le(&mut cursor)?;
        if magic != IVFFLAT_MAGIC {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Invalid IVFFLAT magic: 0x{:08X}", magic),
            ));
        }
        let version = read_u32_le(&mut cursor)?;
        if version != IVFFLAT_VERSION {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Unsupported IVFFLAT version: {}", version),
            ));
        }

        let d = validate_positive_i32(read_i32_le(&mut cursor)?, "d")? as usize;
        let nlist = validate_positive_i32(read_i32_le(&mut cursor)?, "nlist")? as usize;
        let metric_code = read_u32_le(&mut cursor)?;
        let metric = MetricType::from_code(metric_code).ok_or_else(|| {
            io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Unknown metric type: {}", metric_code),
            )
        })?;
        let total_vectors = read_i64_le(&mut cursor)?;
        let flags = read_u32_le(&mut cursor)?;
        let mut reserved = [0u8; 32];
        cursor.read_exact(&mut reserved)?;
        validate_reserved_zero(&reserved, "IVFFLAT")?;
        let unknown_flags = flags & !SUPPORTED_FLAGS;
        if unknown_flags != 0 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Unsupported IVFFLAT flags: 0x{:08X}", unknown_flags),
            ));
        }
        if flags & REQUIRED_FLAGS != REQUIRED_FLAGS {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "IVFFLAT v1 requires delta IDs",
            ));
        }

        Ok(Self {
            reader,
            d,
            nlist,
            metric,
            total_vectors,
            quantizer_centroids: Vec::new(),
            list_offsets: Vec::new(),
            list_counts: Vec::new(),
            list_id_bytes_lens: Vec::new(),
            delta_ids: true,
            loaded: false,
        })
    }

    pub fn ensure_loaded(&mut self) -> io::Result<()> {
        if self.loaded {
            return Ok(());
        }

        let mut cursor = PreadCursor::new(&mut self.reader, IVFFLAT_HEADER_SIZE as u64);
        self.quantizer_centroids =
            read_f32_vec(&mut cursor, checked_section_size(self.nlist, self.d)?)?;
        self.list_offsets = vec![0; self.nlist];
        self.list_counts = vec![0; self.nlist];
        self.list_id_bytes_lens = vec![0; self.nlist];
        for list_id in 0..self.nlist {
            self.list_offsets[list_id] = read_i64_le(&mut cursor)?;
            let count = read_i32_le(&mut cursor)?;
            if count < 0 {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("negative list count {} at list {}", count, list_id),
                ));
            }
            self.list_counts[list_id] = count;
            let id_bytes_len = read_i32_le(&mut cursor)?;
            if id_bytes_len < 0 {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("negative id_bytes_len {} at list {}", id_bytes_len, list_id),
                ));
            }
            self.list_id_bytes_lens[list_id] = id_bytes_len;
        }

        self.loaded = true;
        Ok(())
    }

    pub fn read_inverted_list(&mut self, list_id: usize) -> io::Result<(Vec<i64>, Vec<f32>)> {
        self.ensure_loaded()?;
        if list_id >= self.nlist {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("list_id {} out of range (nlist={})", list_id, self.nlist),
            ));
        }
        let count = self.list_counts[list_id] as usize;
        if count == 0 {
            return Ok((Vec::new(), Vec::new()));
        }

        let offset = checked_list_offset(self.list_offsets[list_id], list_id)?;
        let vector_bytes = checked_list_bytes(count, self.d * 4)?;
        if self.delta_ids {
            let id_bytes_len = self.list_id_bytes_lens[list_id] as usize;
            let payload_len = 12usize
                .checked_add(id_bytes_len)
                .and_then(|len| len.checked_add(vector_bytes))
                .ok_or_else(|| {
                    io::Error::new(io::ErrorKind::InvalidData, "IVF-FLAT list payload overflow")
                })?;
            let mut payload = vec![0u8; payload_len];
            self.reader
                .pread(&mut [ReadRequest::new(offset, &mut payload)])?;
            let base_id = i64::from_le_bytes(payload[0..8].try_into().unwrap());
            let encoded_len = i32::from_le_bytes(payload[8..12].try_into().unwrap());
            if encoded_len < 0 || encoded_len as usize != id_bytes_len {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "IVF-FLAT id_bytes_len mismatch",
                ));
            }
            let ids = decode_delta_varint_ids(base_id, &payload[12..12 + id_bytes_len], count)?;
            let vectors = bytes_to_f32_vec(&payload[12 + id_bytes_len..])?;
            Ok((ids, vectors))
        } else {
            Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "IVF-FLAT reader only supports delta IDs",
            ))
        }
    }

    pub fn search(
        &mut self,
        query: &[f32],
        k: usize,
        nprobe: usize,
    ) -> io::Result<(Vec<i64>, Vec<f32>)> {
        self.search_with_filter(query, k, nprobe, None)
    }

    pub fn search_with_filter(
        &mut self,
        query: &[f32],
        k: usize,
        nprobe: usize,
        filter: Option<&dyn RowIdFilter>,
    ) -> io::Result<(Vec<i64>, Vec<f32>)> {
        self.ensure_loaded()?;
        if query.len() != self.d {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!(
                    "query length {} does not match index dimension {}",
                    query.len(),
                    self.d
                ),
            ));
        }
        if k == 0 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "k must be greater than 0",
            ));
        }
        if nprobe == 0 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "nprobe must be greater than 0",
            ));
        }

        let mut q = query.to_vec();
        if self.metric == MetricType::Cosine {
            fvec_normalize(&mut q);
        }

        let (probe_indices, _) =
            kmeans::find_topk(&q, &self.quantizer_centroids, self.nlist, self.d, nprobe);
        let mut heap = ReaderTopKHeap::new(k);

        for list_id in probe_indices {
            let (ids, vectors) = self.read_inverted_list(list_id)?;
            for (local_idx, &id) in ids.iter().enumerate() {
                if let Some(f) = filter {
                    if !f.contains(id) {
                        continue;
                    }
                }
                let vector = &vectors[local_idx * self.d..(local_idx + 1) * self.d];
                heap.push(fvec_distance(&q, vector, self.metric), id);
            }
        }

        let sorted = heap.into_sorted();
        let mut labels: Vec<i64> = sorted.iter().map(|&(_, id)| id).collect();
        let mut distances: Vec<f32> = sorted.iter().map(|&(dist, _)| dist).collect();
        labels.resize(k, -1);
        distances.resize(k, f32::MAX);
        Ok((labels, distances))
    }

    pub fn search_with_roaring_filter(
        &mut self,
        query: &[f32],
        k: usize,
        nprobe: usize,
        roaring_filter_bytes: &[u8],
    ) -> io::Result<(Vec<i64>, Vec<f32>)> {
        let filter = decode_roaring_filter(roaring_filter_bytes)?;
        self.search_with_filter(query, k, nprobe, Some(&filter))
    }
}

/// Batch search for IVF-FLAT readers. Each unique probed list is read once and
/// scanned for all queries that selected it.
pub fn search_batch_ivfflat_reader<R: SeekRead>(
    reader: &mut IVFFlatIndexReader<R>,
    queries: &[f32],
    nq: usize,
    k: usize,
    nprobe: usize,
) -> io::Result<(Vec<i64>, Vec<f32>)> {
    search_batch_ivfflat_reader_filter(reader, queries, nq, k, nprobe, None)
}

pub fn search_batch_ivfflat_reader_filter<R: SeekRead>(
    reader: &mut IVFFlatIndexReader<R>,
    queries: &[f32],
    nq: usize,
    k: usize,
    nprobe: usize,
    filter: Option<&dyn RowIdFilter>,
) -> io::Result<(Vec<i64>, Vec<f32>)> {
    reader.ensure_loaded()?;
    let d = reader.d;
    if nq == 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "nq must be greater than 0",
        ));
    }
    let expected_query_len = nq.checked_mul(d).ok_or_else(|| {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            "nq * dimension overflows usize",
        )
    })?;
    if queries.len() != expected_query_len {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!(
                "queries length {} does not match nq * dimension {}",
                queries.len(),
                expected_query_len
            ),
        ));
    }
    if k == 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "k must be greater than 0",
        ));
    }
    if nprobe == 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "nprobe must be greater than 0",
        ));
    }

    let mut processed = queries[..expected_query_len].to_vec();
    if reader.metric == MetricType::Cosine {
        for qi in 0..nq {
            fvec_normalize(&mut processed[qi * d..(qi + 1) * d]);
        }
    }

    let (all_probe_indices, _) = kmeans::find_topk_batch(
        &processed,
        nq,
        &reader.quantizer_centroids,
        reader.nlist,
        d,
        nprobe,
    );

    let mut list_to_queries = vec![Vec::new(); reader.nlist];
    let mut unique_lists = Vec::new();
    for (qi, probe_indices) in all_probe_indices.iter().enumerate() {
        for &list_id in probe_indices {
            if list_to_queries[list_id].is_empty() {
                unique_lists.push(list_id);
            }
            list_to_queries[list_id].push(qi);
        }
    }

    let mut heaps: Vec<ReaderTopKHeap> = (0..nq).map(|_| ReaderTopKHeap::new(k)).collect();
    for list_id in unique_lists {
        let count = reader.list_counts[list_id] as usize;
        if count == 0 {
            continue;
        }
        let (ids, vectors) = reader.read_inverted_list(list_id)?;
        for &qi in &list_to_queries[list_id] {
            let query = &processed[qi * d..(qi + 1) * d];
            for (local_idx, &id) in ids.iter().enumerate() {
                if let Some(f) = filter {
                    if !f.contains(id) {
                        continue;
                    }
                }
                let vector = &vectors[local_idx * d..(local_idx + 1) * d];
                heaps[qi].push(fvec_distance(query, vector, reader.metric), id);
            }
        }
    }

    let mut result_ids = vec![-1i64; nq * k];
    let mut result_dists = vec![f32::MAX; nq * k];
    for qi in 0..nq {
        let sorted = std::mem::replace(&mut heaps[qi], ReaderTopKHeap::new(0)).into_sorted();
        let base = qi * k;
        for (i, &(dist, id)) in sorted.iter().enumerate() {
            result_ids[base + i] = id;
            result_dists[base + i] = dist;
        }
    }

    Ok((result_ids, result_dists))
}

pub fn search_batch_ivfflat_reader_roaring_filter<R: SeekRead>(
    reader: &mut IVFFlatIndexReader<R>,
    queries: &[f32],
    nq: usize,
    k: usize,
    nprobe: usize,
    roaring_filter_bytes: &[u8],
) -> io::Result<(Vec<i64>, Vec<f32>)> {
    let filter = decode_roaring_filter(roaring_filter_bytes)?;
    search_batch_ivfflat_reader_filter(reader, queries, nq, k, nprobe, Some(&filter))
}

struct ReaderTopKHeap {
    k: usize,
    data: Vec<(f32, i64)>,
}

impl ReaderTopKHeap {
    fn new(k: usize) -> Self {
        Self {
            k,
            data: Vec::with_capacity(k),
        }
    }

    fn push(&mut self, dist: f32, id: i64) {
        if self.k == 0 {
            return;
        }
        if self.data.len() < self.k {
            self.data.push((dist, id));
            return;
        }
        if let Some((worst_idx, _)) = self
            .data
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.0.partial_cmp(&b.0).unwrap())
        {
            if dist < self.data[worst_idx].0 {
                self.data[worst_idx] = (dist, id);
            }
        }
    }

    fn into_sorted(mut self) -> Vec<(f32, i64)> {
        self.data.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
        self.data
    }
}

fn write_u32_le(out: &mut dyn SeekWrite, v: u32) -> io::Result<()> {
    out.write_all(&v.to_le_bytes())
}

fn write_i32_le(out: &mut dyn SeekWrite, v: i32) -> io::Result<()> {
    out.write_all(&v.to_le_bytes())
}

fn write_i64_le(out: &mut dyn SeekWrite, v: i64) -> io::Result<()> {
    out.write_all(&v.to_le_bytes())
}

fn write_f32_slice(out: &mut dyn SeekWrite, data: &[f32]) -> io::Result<()> {
    let bytes: Vec<u8> = data.iter().flat_map(|f| f.to_le_bytes()).collect();
    out.write_all(&bytes)
}

fn read_u32_le<R: SeekRead + ?Sized>(reader: &mut PreadCursor<'_, R>) -> io::Result<u32> {
    let mut buf = [0u8; 4];
    reader.read_exact(&mut buf)?;
    Ok(u32::from_le_bytes(buf))
}

fn read_i32_le<R: SeekRead + ?Sized>(reader: &mut PreadCursor<'_, R>) -> io::Result<i32> {
    let mut buf = [0u8; 4];
    reader.read_exact(&mut buf)?;
    Ok(i32::from_le_bytes(buf))
}

fn read_i64_le<R: SeekRead + ?Sized>(reader: &mut PreadCursor<'_, R>) -> io::Result<i64> {
    let mut buf = [0u8; 8];
    reader.read_exact(&mut buf)?;
    Ok(i64::from_le_bytes(buf))
}

fn validate_positive_i32(val: i32, field: &str) -> io::Result<i32> {
    if val <= 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("invalid header field {}: {} (must be positive)", field, val),
        ));
    }
    Ok(val)
}

fn validate_index_shape(index: &IVFFlatIndex) -> io::Result<()> {
    if index.d == 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "dimension must be greater than 0",
        ));
    }
    if index.nlist == 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "nlist must be greater than 0",
        ));
    }
    if index.ids.len() != index.nlist || index.vectors.len() != index.nlist {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "IVF-FLAT list storage does not match nlist",
        ));
    }
    let centroid_len = checked_section_size(index.nlist, index.d)?;
    if index.quantizer_centroids.len() != centroid_len {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!(
                "centroid length {} does not match nlist*d {}",
                index.quantizer_centroids.len(),
                centroid_len
            ),
        ));
    }
    for list_id in 0..index.nlist {
        let expected_vector_len =
            index.ids[list_id]
                .len()
                .checked_mul(index.d)
                .ok_or_else(|| {
                    io::Error::new(
                        io::ErrorKind::InvalidInput,
                        "IVF-FLAT vector length overflow",
                    )
                })?;
        if index.vectors[list_id].len() != expected_vector_len {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!(
                    "list {} vector length {} does not match ids*d {}",
                    list_id,
                    index.vectors[list_id].len(),
                    expected_vector_len
                ),
            ));
        }
    }
    Ok(())
}

fn usize_to_i32(value: usize, field: &str) -> io::Result<i32> {
    if value > i32::MAX as usize {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("{} exceeds i32 length limit: {}", field, value),
        ));
    }
    Ok(value as i32)
}

fn usize_to_i64(value: usize, field: &str) -> io::Result<i64> {
    if value > i64::MAX as usize {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("{} exceeds i64 length limit: {}", field, value),
        ));
    }
    Ok(value as i64)
}

fn u64_to_i64(value: u64, field: &str) -> io::Result<i64> {
    if value > i64::MAX as u64 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("{} exceeds i64 offset limit: {}", field, value),
        ));
    }
    Ok(value as i64)
}

const MAX_SECTION_ELEMENTS: usize = 1 << 30;

fn checked_section_size(a: usize, b: usize) -> io::Result<usize> {
    let result = a.checked_mul(b).ok_or_else(|| {
        io::Error::new(
            io::ErrorKind::InvalidData,
            "section size overflow in IVF-FLAT header",
        )
    })?;
    if result > MAX_SECTION_ELEMENTS {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!(
                "section size {} exceeds maximum {}",
                result, MAX_SECTION_ELEMENTS
            ),
        ));
    }
    Ok(result)
}

fn checked_list_offset(offset: i64, list_id: usize) -> io::Result<u64> {
    if offset < 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("negative list offset {} at list {}", offset, list_id),
        ));
    }
    Ok(offset as u64)
}

fn checked_list_bytes(count: usize, bytes_per_entry: usize) -> io::Result<usize> {
    count.checked_mul(bytes_per_entry).ok_or_else(|| {
        io::Error::new(
            io::ErrorKind::InvalidData,
            "IVF-FLAT list byte size overflow",
        )
    })
}

fn read_f32_vec<R: SeekRead + ?Sized>(
    reader: &mut PreadCursor<'_, R>,
    count: usize,
) -> io::Result<Vec<f32>> {
    let mut buf = vec![0u8; count * 4];
    reader.read_exact(&mut buf)?;
    bytes_to_f32_vec(&buf)
}

fn bytes_to_f32_vec(bytes: &[u8]) -> io::Result<Vec<f32>> {
    if !bytes.len().is_multiple_of(4) {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "f32 byte section is not 4-byte aligned",
        ));
    }
    Ok(bytes
        .chunks_exact(4)
        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
        .collect())
}

fn encode_varint(mut val: u64, buf: &mut Vec<u8>) {
    while val >= 0x80 {
        buf.push((val as u8) | 0x80);
        val >>= 7;
    }
    buf.push(val as u8);
}

fn decode_varint(buf: &[u8], pos: &mut usize) -> io::Result<u64> {
    let mut val = 0u64;
    let mut shift = 0u32;
    loop {
        if *pos >= buf.len() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "truncated varint",
            ));
        }
        let b = buf[*pos] as u64;
        *pos += 1;
        let payload = b & 0x7F;
        if shift == 63 && payload > 1 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "varint exceeds u64 range",
            ));
        }
        val |= payload << shift;
        if b & 0x80 == 0 {
            break;
        }
        shift += 7;
        if shift > 63 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "varint exceeds 64 bits",
            ));
        }
    }
    Ok(val)
}

fn encode_delta_varint_ids(ids: &[i64]) -> (i64, Vec<u8>) {
    if ids.is_empty() {
        return (0, Vec::new());
    }
    let base = ids[0];
    let mut buf = Vec::with_capacity(ids.len() * 2);
    let mut prev = base;
    for &id in ids {
        let delta = (id as u64).wrapping_sub(prev as u64);
        encode_varint(delta, &mut buf);
        prev = id;
    }
    (base, buf)
}

fn decode_delta_varint_ids(base: i64, buf: &[u8], count: usize) -> io::Result<Vec<i64>> {
    let mut ids = Vec::with_capacity(count);
    let mut pos = 0;
    let mut current = base as u64;
    let mut prev_signed = base;
    for _ in 0..count {
        let delta = decode_varint(buf, &mut pos)?;
        current = current.wrapping_add(delta);
        let id = current as i64;
        if id < prev_signed {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "decoded ID sequence is not monotonically non-decreasing",
            ));
        }
        prev_signed = id;
        ids.push(id);
    }
    Ok(ids)
}

fn decode_roaring_filter(bytes: &[u8]) -> io::Result<RoaringTreemap> {
    RoaringTreemap::deserialize_from(bytes).map_err(|e| {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("invalid RoaringTreemap filter: {}", e),
        )
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::distance::MetricType;
    use crate::io::PosWriter;
    use crate::ivfflat::IVFFlatIndex;
    use std::io::Cursor;

    #[test]
    fn test_ivfflat_write_read_search_roundtrip() {
        let d = 4;
        let nlist = 4;
        let n = 128;
        let data: Vec<f32> = (0..n)
            .flat_map(|i| {
                let cluster = (i % nlist) as f32 * 100.0;
                [cluster + i as f32 * 0.01, 1.0, 2.0, 3.0]
            })
            .collect();
        let ids: Vec<i64> = (1000..1000 + n as i64).collect();

        let mut index = IVFFlatIndex::new(d, nlist, MetricType::L2);
        index.train(&data, n);
        index.add(&data, &ids, n);

        let mut expected_distances = vec![0.0; 5];
        let mut expected_labels = vec![0; 5];
        index.search(
            &data[7 * d..8 * d],
            1,
            5,
            nlist,
            &mut expected_distances,
            &mut expected_labels,
        );

        let mut buf = Vec::new();
        let mut writer = PosWriter::new(&mut buf);
        write_ivfflat_index(&index, &mut writer).unwrap();

        let mut reader = IVFFlatIndexReader::open(Cursor::new(buf)).unwrap();
        let (labels, distances) = reader.search(&data[7 * d..8 * d], 5, nlist).unwrap();

        assert_eq!(labels, expected_labels);
        assert_eq!(distances, expected_distances);
    }

    #[test]
    fn test_ivfflat_reader_search_with_filter() {
        use std::collections::HashSet;

        let d = 2;
        let nlist = 1;
        let data = vec![0.0, 0.0, 0.1, 0.0, 10.0, 10.0];
        let ids = vec![10, 11, 12];

        let mut index = IVFFlatIndex::new(d, nlist, MetricType::L2);
        index.train(&data, 3);
        index.add(&data, &ids, 3);

        let mut buf = Vec::new();
        let mut writer = PosWriter::new(&mut buf);
        write_ivfflat_index(&index, &mut writer).unwrap();

        let filter: HashSet<i64> = [12].into_iter().collect();
        let mut reader = IVFFlatIndexReader::open(Cursor::new(buf)).unwrap();
        let (labels, distances) = reader
            .search_with_filter(&[0.0, 0.0], 2, 1, Some(&filter))
            .unwrap();

        assert_eq!(labels, vec![12, -1]);
        assert_eq!(distances[0], 200.0);
        assert_eq!(distances[1], f32::MAX);
    }

    #[test]
    fn test_ivfflat_reader_search_with_roaring_filter_bytes() {
        let d = 2;
        let nlist = 1;
        let data = vec![0.0, 0.0, 0.1, 0.0, 10.0, 10.0];
        let ids = vec![10, 11, 12];

        let mut index = IVFFlatIndex::new(d, nlist, MetricType::L2);
        index.train(&data, 3);
        index.add(&data, &ids, 3);

        let mut buf = Vec::new();
        let mut writer = PosWriter::new(&mut buf);
        write_ivfflat_index(&index, &mut writer).unwrap();

        let mut allowed = RoaringTreemap::new();
        allowed.insert(12);
        let mut filter_bytes = Vec::new();
        allowed.serialize_into(&mut filter_bytes).unwrap();

        let mut reader = IVFFlatIndexReader::open(Cursor::new(buf)).unwrap();
        let (labels, distances) = reader
            .search_with_roaring_filter(&[0.0, 0.0], 2, 1, &filter_bytes)
            .unwrap();

        assert_eq!(labels, vec![12, -1]);
        assert_eq!(distances[0], 200.0);
        assert_eq!(distances[1], f32::MAX);
    }

    #[test]
    fn test_ivfflat_batch_reader_matches_single_reader_search() {
        let d = 4;
        let nlist = 4;
        let n = 128;
        let data: Vec<f32> = (0..n)
            .flat_map(|i| {
                let cluster = (i % nlist) as f32 * 100.0;
                [cluster + i as f32 * 0.01, 1.0, 2.0, 3.0]
            })
            .collect();
        let ids: Vec<i64> = (1000..1000 + n as i64).collect();

        let mut index = IVFFlatIndex::new(d, nlist, MetricType::L2);
        index.train(&data, n);
        index.add(&data, &ids, n);

        let mut buf = Vec::new();
        let mut writer = PosWriter::new(&mut buf);
        write_ivfflat_index(&index, &mut writer).unwrap();

        let queries = [&data[7 * d..8 * d], &data[63 * d..64 * d]].concat();
        let k = 5;
        let nprobe = 3;
        let mut batch_reader = IVFFlatIndexReader::open(Cursor::new(buf.clone())).unwrap();
        let (batch_labels, batch_distances) =
            search_batch_ivfflat_reader(&mut batch_reader, &queries, 2, k, nprobe).unwrap();

        for qi in 0..2 {
            let mut single_reader = IVFFlatIndexReader::open(Cursor::new(buf.clone())).unwrap();
            let query = &queries[qi * d..(qi + 1) * d];
            let (single_labels, single_distances) = single_reader.search(query, k, nprobe).unwrap();
            assert_eq!(&batch_labels[qi * k..(qi + 1) * k], single_labels);
            assert_eq!(&batch_distances[qi * k..(qi + 1) * k], single_distances);
        }
    }

    #[test]
    fn test_ivfflat_batch_reader_search_with_roaring_filter_bytes() {
        let d = 2;
        let nlist = 1;
        let data = vec![0.0, 0.0, 0.1, 0.0, 10.0, 10.0];
        let ids = vec![10, 11, 12];

        let mut index = IVFFlatIndex::new(d, nlist, MetricType::L2);
        index.train(&data, 3);
        index.add(&data, &ids, 3);

        let mut buf = Vec::new();
        let mut writer = PosWriter::new(&mut buf);
        write_ivfflat_index(&index, &mut writer).unwrap();

        let mut allowed = RoaringTreemap::new();
        allowed.insert(12);
        let mut filter_bytes = Vec::new();
        allowed.serialize_into(&mut filter_bytes).unwrap();

        let mut reader = IVFFlatIndexReader::open(Cursor::new(buf)).unwrap();
        let queries = vec![0.0, 0.0, 10.0, 10.0];
        let (labels, distances) = search_batch_ivfflat_reader_roaring_filter(
            &mut reader,
            &queries,
            2,
            2,
            1,
            &filter_bytes,
        )
        .unwrap();

        assert_eq!(labels, vec![12, -1, 12, -1]);
        assert_eq!(distances, vec![200.0, f32::MAX, 0.0, f32::MAX]);
    }

    #[test]
    fn test_ivfflat_reader_validates_inputs() {
        let d = 2;
        let nlist = 1;
        let data = vec![0.0, 0.0, 1.0, 1.0];
        let ids = vec![1, 2];

        let mut index = IVFFlatIndex::new(d, nlist, MetricType::L2);
        index.train(&data, 2);
        index.add(&data, &ids, 2);

        let mut buf = Vec::new();
        let mut writer = PosWriter::new(&mut buf);
        write_ivfflat_index(&index, &mut writer).unwrap();

        let mut reader = IVFFlatIndexReader::open(Cursor::new(buf.clone())).unwrap();
        assert!(reader.search(&[0.0], 1, 1).is_err());

        let mut reader = IVFFlatIndexReader::open(Cursor::new(buf.clone())).unwrap();
        assert!(reader.search(&[0.0, 0.0], 0, 1).is_err());

        let mut reader = IVFFlatIndexReader::open(Cursor::new(buf)).unwrap();
        assert!(reader.search(&[0.0, 0.0], 1, 0).is_err());
    }

    #[test]
    fn test_ivfflat_writer_validates_shape_before_writing() {
        let mut index = IVFFlatIndex::new(2, 1, MetricType::L2);
        index.quantizer_centroids = vec![0.0, 0.0];
        index.ids[0] = vec![1, 2];
        index.vectors[0] = vec![0.0, 0.0];

        let mut buf = Vec::new();
        let mut writer = PosWriter::new(&mut buf);
        let err = write_ivfflat_index(&index, &mut writer).unwrap_err();

        assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
        assert!(err.to_string().contains("vector length"));
    }

    #[test]
    fn test_ivfflat_reader_rejects_bad_magic() {
        let mut buf = vec![0u8; IVFFLAT_HEADER_SIZE];
        buf[0..4].copy_from_slice(&0x12345678u32.to_le_bytes());

        let err = match IVFFlatIndexReader::open(Cursor::new(buf)) {
            Ok(_) => panic!("bad magic should be rejected"),
            Err(err) => err,
        };
        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
    }

    #[test]
    fn test_ivfflat_reader_rejects_missing_required_flags() {
        let mut buf = vec![0u8; IVFFLAT_HEADER_SIZE];
        buf[0..4].copy_from_slice(&IVFFLAT_MAGIC.to_le_bytes());
        buf[4..8].copy_from_slice(&IVFFLAT_VERSION.to_le_bytes());
        buf[8..12].copy_from_slice(&2i32.to_le_bytes());
        buf[12..16].copy_from_slice(&1i32.to_le_bytes());
        buf[16..20].copy_from_slice(&(MetricType::L2 as u32).to_le_bytes());
        buf[20..28].copy_from_slice(&0i64.to_le_bytes());
        buf[28..32].copy_from_slice(&0u32.to_le_bytes());

        let err = match IVFFlatIndexReader::open(Cursor::new(buf)) {
            Ok(_) => panic!("missing required flags should be rejected"),
            Err(err) => err,
        };
        assert!(err.to_string().contains("requires delta IDs"));
    }

    #[test]
    fn test_ivfflat_reader_rejects_unknown_flags() {
        let mut buf = vec![0u8; IVFFLAT_HEADER_SIZE];
        buf[0..4].copy_from_slice(&IVFFLAT_MAGIC.to_le_bytes());
        buf[4..8].copy_from_slice(&IVFFLAT_VERSION.to_le_bytes());
        buf[8..12].copy_from_slice(&2i32.to_le_bytes());
        buf[12..16].copy_from_slice(&1i32.to_le_bytes());
        buf[16..20].copy_from_slice(&(MetricType::L2 as u32).to_le_bytes());
        buf[20..28].copy_from_slice(&0i64.to_le_bytes());
        buf[28..32].copy_from_slice(&(REQUIRED_FLAGS | (1 << 31)).to_le_bytes());

        let err = match IVFFlatIndexReader::open(Cursor::new(buf)) {
            Ok(_) => panic!("unknown flags should be rejected"),
            Err(err) => err,
        };
        assert!(err.to_string().contains("Unsupported IVFFLAT flags"));
    }

    #[test]
    fn test_ivfflat_reader_rejects_nonzero_reserved_bytes() {
        let mut buf = vec![0u8; IVFFLAT_HEADER_SIZE];
        buf[0..4].copy_from_slice(&IVFFLAT_MAGIC.to_le_bytes());
        buf[4..8].copy_from_slice(&IVFFLAT_VERSION.to_le_bytes());
        buf[8..12].copy_from_slice(&2i32.to_le_bytes());
        buf[12..16].copy_from_slice(&1i32.to_le_bytes());
        buf[16..20].copy_from_slice(&(MetricType::L2 as u32).to_le_bytes());
        buf[20..28].copy_from_slice(&0i64.to_le_bytes());
        buf[28..32].copy_from_slice(&REQUIRED_FLAGS.to_le_bytes());
        buf[32] = 1;

        let err = match IVFFlatIndexReader::open(Cursor::new(buf)) {
            Ok(_) => panic!("non-zero reserved bytes should be rejected"),
            Err(err) => err,
        };
        assert!(err.to_string().contains("reserved bytes must be zero"));
    }
}
