qdrant-edge 0.7.2

use std::path::PathBuf;

use crate::common::bitvec::{BitSlice, BitVec};
use crate::common::types::PointOffsetType;
use crate::common::universal_io::{UniversalRead, UniversalWrite};

use super::buffered_dynamic_flags::BufferedDynamicFlags;
use super::dynamic_stored_flags::DynamicStoredFlags;
use crate::segment::common::Flusher;
use crate::segment::common::operation_error::OperationResult;

/// A buffered, growable, and persistent bitslice with a separate in-memory bitvec.
///
/// Use [`RoaringFlags`][1] if you need a reference to a bitmap.
///
/// Changes are buffered until explicitly flushed.
///
/// [1]: super::roaring_flags::RoaringFlags
#[derive(Debug)]
pub struct BitvecFlags<S: UniversalRead> {
    /// Buffered persisted flags.
    storage: BufferedDynamicFlags<S>,

    /// In-memory bitvec of true and false flags.
    bitvec: BitVec,

    /// Total length of the flags, including the trailing ones which have been set to false
    len: usize,
}

impl<S> BitvecFlags<S>
where
    S: UniversalWrite + Send + 'static,
    S::Fs: Send + Sync + 'static,
{
    pub fn new(fs: S::Fs, dynamic_flags: DynamicStoredFlags<S>) -> OperationResult<Self> {
        // load flags into memory
        let bitvec = BitVec::from_bitslice(&*dynamic_flags.get_bitslice()?);

        if let Err(err) = dynamic_flags.clear_cache() {
            log::warn!("Failed to clear bitslice cache: {err}");
        }

        Ok(Self {
            len: dynamic_flags.len(),
            storage: BufferedDynamicFlags::new(fs, dynamic_flags),
            bitvec,
        })
    }

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    pub fn get_bitslice(&self) -> &BitSlice {
        &self.bitvec
    }

    pub fn get(&self, index: PointOffsetType) -> bool {
        self.bitvec.get(index as usize).is_some_and(|bit| *bit)
    }

    pub fn iter_trues(&self) -> impl Iterator<Item = PointOffsetType> {
        self.bitvec
            .iter_ones()
            .map(|index| index as PointOffsetType)
    }

    pub fn iter_falses(&self) -> impl Iterator<Item = PointOffsetType> {
        self.bitvec
            .iter_zeros()
            .map(|index| index as PointOffsetType)
    }

    #[inline]
    pub fn count_trues(&self) -> usize {
        self.bitvec.count_ones()
    }

    #[inline]
    pub fn count_falses(&self) -> usize {
        self.bitvec.count_zeros()
    }

    /// Set the value of a flag at the given index, grows the bitvec if needed.
    /// Returns the previous value of the flag.
    pub fn set(&mut self, index: PointOffsetType, value: bool) -> bool {
        // queue write in buffer
        self.storage.buffer_set(index, value);

        // update length if needed
        let index_usize = index as usize;
        if index_usize >= self.len {
            self.len = index_usize + 1;
            self.bitvec.resize(self.len, false);
        }

        // update bitmap
        self.bitvec.replace(index_usize, value)
    }

    pub fn clear_cache(&self) -> OperationResult<()> {
        let Self {
            storage,
            bitvec: _,
            len: _,
        } = self;
        storage.clear_cache()?;
        Ok(())
    }

    pub fn files(&self) -> Vec<PathBuf> {
        self.storage.files()
    }

    pub fn flusher(&self) -> Flusher {
        self.storage.flusher()
    }
}

#[allow(clippy::default_constructed_unit_structs)]
#[duplicate::duplicate_item(
    tests_mod       S               Fs              cfg_predicate;
    [tests_mmap]    [MmapFile]      [MmapFs]        [cfg(all())];
    [tests_uring]   [IoUringFile]   [IoUringFs]     [cfg(target_os = "linux")];
)]
#[cfg_predicate]
#[cfg(test)]
mod tests_mod {
    use crate::common::types::PointOffsetType;
    #[cfg_predicate]
    use crate::common::universal_io::{Fs, S};

    use crate::segment::common::flags::bitvec_flags::BitvecFlags;
    use crate::segment::common::flags::dynamic_stored_flags::DynamicStoredFlags;

    #[test]
    fn test_roaring_flags_consistency_after_persistence() {
        let dir = tempfile::Builder::new()
            .prefix("roaring_flags_consistency")
            .tempdir()
            .unwrap();

        // Create and update flags
        {
            let mmap_flags =
                DynamicStoredFlags::<S>::open(&Fs::default(), dir.path(), false).unwrap();
            let mut bitvec_flags = BitvecFlags::new(Fs::default(), mmap_flags).unwrap();

            // Set various flags - we'll set up to index 19 to have a length of 20
            for i in 16..20 {
                bitvec_flags.set(i, false); // Ensure we have length 20
            }
            bitvec_flags.set(0, true);
            bitvec_flags.set(5, true);
            bitvec_flags.set(10, true);
            bitvec_flags.set(15, true);
            bitvec_flags.set(7, false); // This should be no-op since default is false

            // Verify iteration consistency after reload
            let iter_trues: Vec<_> = bitvec_flags.iter_trues().collect();

            // Verify expected values
            assert_eq!(iter_trues, vec![0, 5, 10, 15]);

            // Verify count consistency
            assert_eq!(bitvec_flags.count_trues(), 4);

            // Flush
            let flusher = bitvec_flags.flusher();
            flusher().unwrap();
        }

        // Verify bitmap consistency after reload
        {
            let mmap_flags =
                DynamicStoredFlags::<S>::open(&Fs::default(), dir.path(), true).unwrap();
            let bitvec_flags = BitvecFlags::new(Fs::default(), mmap_flags).unwrap();

            // Verify iteration consistency after reload
            let iter_trues: Vec<_> = bitvec_flags.iter_trues().collect();

            // Verify expected values
            assert_eq!(iter_trues, vec![0, 5, 10, 15]);

            // Verify count consistency
            assert_eq!(bitvec_flags.count_trues(), 4);
            assert_eq!(
                bitvec_flags.count_falses(),
                bitvec_flags.len() - bitvec_flags.count_trues()
            );

            // Verify iteration covers all indices
            let all_trues: Vec<_> = bitvec_flags.iter_trues().collect();
            let all_falses: Vec<_> = bitvec_flags.iter_falses().collect();
            let mut all_indices = all_trues;
            all_indices.extend(all_falses);
            all_indices.sort();

            let expected_all: Vec<_> = (0..bitvec_flags.len() as PointOffsetType).collect();
            assert_eq!(all_indices, expected_all);
        }
    }
}