qdrant-edge 0.7.1

use std::path::{Path, PathBuf};

use crate::common::bitvec::BitSlice;
use crate::gridstore::Blob;

use super::bool_index::BoolIndex;
use super::bool_index::immutable_bool_index::ImmutableBoolIndex;
use super::bool_index::mutable_bool_index::MutableBoolIndex;
use super::geo_index::{GeoMapIndexGridstoreBuilder, GeoMapIndexMmapBuilder};
use super::map_index::{MapIndex, MapIndexGridstoreBuilder, MapIndexKey, MapIndexMmapBuilder};
use super::null_index::{ImmutableNullIndex, NullIndex};
use super::numeric_index::{
    Encodable, NumericIndexGridstoreBuilder, NumericIndexIntoInnerValue, NumericIndexMmapBuilder,
};
use super::stored_point_to_values::StoredValue;
use super::{FieldIndexBuilder, ValueIndexer};
use crate::segment::common::operation_error::{OperationError, OperationResult};
use crate::segment::data_types::index::TextIndexParams;
use crate::segment::id_tracker::{IdTrackerEnum, IdTrackerRead};
use crate::segment::index::field_index::FieldIndex;
use crate::segment::index::field_index::full_text_index::FullTextIndex;
use crate::segment::index::field_index::geo_index::GeoMapIndex;
use crate::segment::index::field_index::null_index::MutableNullIndex;
use crate::segment::index::field_index::numeric_index::NumericIndex;
use crate::segment::index::field_index::numeric_point::Numericable;
use crate::segment::index::payload_config::{FullPayloadIndexType, IndexMutability, PayloadIndexType};
use crate::segment::json_path::JsonPath;
use crate::segment::types::{PayloadFieldSchema, PayloadSchemaParams};

/// Selects index and index builder types based on field type.
#[derive(Copy, Clone)]
pub enum IndexSelector<'a> {
    /// On disk or in-memory index on mmaps, non-appendable
    Mmap(IndexSelectorMmap<'a>),
    /// In-memory index on gridstore, appendable
    Gridstore(IndexSelectorGridstore<'a>),
}

#[derive(Copy, Clone)]
pub struct IndexSelectorMmap<'a> {
    pub dir: &'a Path,
    pub is_on_disk: bool,
}

#[derive(Copy, Clone)]
pub struct IndexSelectorGridstore<'a> {
    pub dir: &'a Path,
}

impl IndexSelector<'_> {
    /// Loads the correct index based on `index_type`.
    #[allow(clippy::too_many_arguments)]
    pub fn new_index_with_type(
        &self,
        field: &JsonPath,
        payload_schema: &PayloadFieldSchema,
        index_type: &FullPayloadIndexType,
        create_if_missing: bool,
        id_tracker: &IdTrackerEnum,
        deleted_points: &BitSlice,
    ) -> OperationResult<Option<FieldIndex>> {
        let index = match (&index_type.index_type, payload_schema.expand().as_ref()) {
            (PayloadIndexType::IntIndex, PayloadSchemaParams::Integer(params)) => {
                // IntIndex only gets created if `range` is true. This will only throw an error if storage is corrupt.
                //
                // Note that `params.range == None` means the index was created without directly specifying these parameters.
                // In those cases it defaults to `true` so we don't need to cover this case.
                if params.range == Some(false) {
                    log::warn!(
                        "Inconsistent payload schema: Int index configured but schema.range is false"
                    );
                }

                self.numeric_new(field, create_if_missing, deleted_points)?
                    .map(FieldIndex::IntIndex)
            }
            (PayloadIndexType::IntMapIndex, PayloadSchemaParams::Integer(params)) => {
                // IntMapIndex only gets created if `lookup` is true. This will only throw an error if storage is corrupt.
                //
                // Note that `params.lookup == None` means the index was created without directly specifying these parameters.
                // In those cases it defaults to `true` so we don't need to cover this case.
                if params.lookup == Some(false) {
                    log::warn!(
                        "Inconsistent payload schema: IntMap index configured but schema.lookup is false",
                    );
                }

                self.map_new(field, create_if_missing, deleted_points)?
                    .map(FieldIndex::IntMapIndex)
            }
            (PayloadIndexType::DatetimeIndex, PayloadSchemaParams::Datetime(_)) => self
                .numeric_new(field, create_if_missing, deleted_points)?
                .map(FieldIndex::DatetimeIndex),

            (PayloadIndexType::KeywordIndex, PayloadSchemaParams::Keyword(_)) => self
                .map_new(field, create_if_missing, deleted_points)?
                .map(FieldIndex::KeywordIndex),

            (PayloadIndexType::FloatIndex, PayloadSchemaParams::Float(_)) => self
                .numeric_new(field, create_if_missing, deleted_points)?
                .map(FieldIndex::FloatIndex),

            (PayloadIndexType::GeoIndex, PayloadSchemaParams::Geo(_)) => self
                .geo_new(field, create_if_missing, deleted_points)?
                .map(FieldIndex::GeoIndex),

            (PayloadIndexType::FullTextIndex, PayloadSchemaParams::Text(params)) => self
                .text_new(field, params.clone(), create_if_missing, deleted_points)?
                .map(FieldIndex::FullTextIndex),

            (PayloadIndexType::BoolIndex, PayloadSchemaParams::Bool(_)) => self
                .bool_new(
                    field,
                    create_if_missing,
                    deleted_points,
                    index_type.mutability,
                )?
                .map(FieldIndex::BoolIndex),

            (PayloadIndexType::UuidIndex, PayloadSchemaParams::Uuid(_)) => self
                .map_new(field, create_if_missing, deleted_points)?
                .map(FieldIndex::UuidMapIndex),

            (PayloadIndexType::UuidMapIndex, PayloadSchemaParams::Uuid(_)) => self
                .map_new(field, create_if_missing, deleted_points)?
                .map(FieldIndex::UuidMapIndex),

            (PayloadIndexType::NullIndex, _) => {
                self.new_null_index(field, create_if_missing, id_tracker, index_type.mutability)?
            }

            // Storage inconsistency. Should never happen.
            (index_type, schema) => {
                return Err(OperationError::service_error(format!(
                    "Payload index storage inconsistent. Schema defines {schema:?} but storage is {index_type:?}"
                )));
            }
        };

        Ok(index)
    }

    /// Selects index type based on field type.
    pub fn new_index(
        &self,
        field: &JsonPath,
        payload_schema: &PayloadFieldSchema,
        create_if_missing: bool,
        deleted_points: &BitSlice,
    ) -> OperationResult<Option<Vec<FieldIndex>>> {
        let indexes = match payload_schema.expand().as_ref() {
            PayloadSchemaParams::Keyword(_) => self
                .map_new(field, create_if_missing, deleted_points)?
                .map(|index| vec![FieldIndex::KeywordIndex(index)]),
            PayloadSchemaParams::Integer(integer_params) => {
                let use_lookup = integer_params.lookup.unwrap_or(true);
                let use_range = integer_params.range.unwrap_or(true);

                let lookup = if use_lookup {
                    match self.map_new(field, create_if_missing, deleted_points)? {
                        Some(index) => Some(FieldIndex::IntMapIndex(index)),
                        None => return Ok(None),
                    }
                } else {
                    None
                };
                let range = if use_range {
                    match self.numeric_new(field, create_if_missing, deleted_points)? {
                        Some(index) => Some(FieldIndex::IntIndex(index)),
                        None => return Ok(None),
                    }
                } else {
                    None
                };

                Some(lookup.into_iter().chain(range).collect())
            }
            PayloadSchemaParams::Float(_) => self
                .numeric_new(field, create_if_missing, deleted_points)?
                .map(|index| vec![FieldIndex::FloatIndex(index)]),
            PayloadSchemaParams::Geo(_) => self
                .geo_new(field, create_if_missing, deleted_points)?
                .map(|index| vec![FieldIndex::GeoIndex(index)]),
            PayloadSchemaParams::Text(text_index_params) => self
                .text_new(
                    field,
                    text_index_params.clone(),
                    create_if_missing,
                    deleted_points,
                )?
                .map(|index| vec![FieldIndex::FullTextIndex(index)]),
            PayloadSchemaParams::Bool(_) => self
                .bool_new(
                    field,
                    create_if_missing,
                    deleted_points,
                    self.default_mutability(),
                )?
                .map(|index| vec![FieldIndex::BoolIndex(index)]),
            PayloadSchemaParams::Datetime(_) => self
                .numeric_new(field, create_if_missing, deleted_points)?
                .map(|index| vec![FieldIndex::DatetimeIndex(index)]),
            PayloadSchemaParams::Uuid(_) => self
                .map_new(field, create_if_missing, deleted_points)?
                .map(|index| vec![FieldIndex::UuidMapIndex(index)]),
        };

        Ok(indexes)
    }

    /// Selects index builder based on field type.
    pub fn index_builder(
        &self,
        field: &JsonPath,
        payload_schema: &PayloadFieldSchema,
        deleted_points: &BitSlice,
    ) -> OperationResult<Vec<FieldIndexBuilder>> {
        let builders = match payload_schema.expand().as_ref() {
            PayloadSchemaParams::Keyword(_) => {
                vec![self.map_builder(
                    field,
                    FieldIndexBuilder::KeywordMmapIndex,
                    FieldIndexBuilder::KeywordGridstoreIndex,
                    deleted_points,
                )]
            }
            PayloadSchemaParams::Integer(integer_params) => {
                let use_lookup = integer_params.lookup.unwrap_or(true);
                let use_range = integer_params.range.unwrap_or(true);

                let lookup = if use_lookup {
                    Some(self.map_builder(
                        field,
                        FieldIndexBuilder::IntMapMmapIndex,
                        FieldIndexBuilder::IntMapGridstoreIndex,
                        deleted_points,
                    ))
                } else {
                    None
                };

                let range = if use_range {
                    Some(self.numeric_builder(
                        field,
                        FieldIndexBuilder::IntMmapIndex,
                        FieldIndexBuilder::IntGridstoreIndex,
                        deleted_points,
                    ))
                } else {
                    None
                };

                lookup.into_iter().chain(range).collect()
            }
            PayloadSchemaParams::Float(_) => {
                vec![self.numeric_builder(
                    field,
                    FieldIndexBuilder::FloatMmapIndex,
                    FieldIndexBuilder::FloatGridstoreIndex,
                    deleted_points,
                )]
            }
            PayloadSchemaParams::Geo(_) => {
                vec![self.geo_builder(
                    field,
                    FieldIndexBuilder::GeoMmapIndex,
                    FieldIndexBuilder::GeoGridstoreIndex,
                    deleted_points,
                )]
            }
            PayloadSchemaParams::Text(text_index_params) => {
                vec![self.text_builder(field, text_index_params.clone(), deleted_points)]
            }
            PayloadSchemaParams::Bool(_) => {
                vec![self.bool_builder(field)?]
            }
            PayloadSchemaParams::Datetime(_) => {
                vec![self.numeric_builder(
                    field,
                    FieldIndexBuilder::DatetimeMmapIndex,
                    FieldIndexBuilder::DatetimeGridstoreIndex,
                    deleted_points,
                )]
            }
            PayloadSchemaParams::Uuid(_) => {
                vec![self.map_builder(
                    field,
                    FieldIndexBuilder::UuidMmapIndex,
                    FieldIndexBuilder::UuidGridstoreIndex,
                    deleted_points,
                )]
            }
        };

        Ok(builders)
    }

    fn map_new<N: MapIndexKey + ?Sized>(
        &self,
        field: &JsonPath,
        create_if_missing: bool,
        deleted_points: &BitSlice,
    ) -> OperationResult<Option<MapIndex<N>>>
    where
        Vec<<N as MapIndexKey>::Owned>: Blob + Send + Sync,
    {
        Ok(match self {
            IndexSelector::Mmap(IndexSelectorMmap { dir, is_on_disk }) => {
                MapIndex::new_mmap(&map_dir(dir, field), *is_on_disk, deleted_points)?
            }
            IndexSelector::Gridstore(IndexSelectorGridstore { dir }) => {
                MapIndex::new_gridstore(map_dir(dir, field), create_if_missing)?
            }
        })
    }

    fn map_builder<N: MapIndexKey + ?Sized>(
        &self,
        field: &JsonPath,
        make_mmap: fn(MapIndexMmapBuilder<N>) -> FieldIndexBuilder,
        make_gridstore: fn(MapIndexGridstoreBuilder<N>) -> FieldIndexBuilder,
        deleted_points: &BitSlice,
    ) -> FieldIndexBuilder
    where
        Vec<<N as MapIndexKey>::Owned>: Blob + Send + Sync,
    {
        match self {
            IndexSelector::Mmap(IndexSelectorMmap { dir, is_on_disk }) => make_mmap(
                MapIndex::builder_mmap(&map_dir(dir, field), *is_on_disk, deleted_points),
            ),
            IndexSelector::Gridstore(IndexSelectorGridstore { dir }) => {
                make_gridstore(MapIndex::builder_gridstore(map_dir(dir, field)))
            }
        }
    }

    fn numeric_new<T: Encodable + Numericable + StoredValue + Send + Sync + Default, P>(
        &self,
        field: &JsonPath,
        create_if_missing: bool,
        deleted_points: &BitSlice,
    ) -> OperationResult<Option<NumericIndex<T, P>>>
    where
        Vec<T>: Blob,
    {
        Ok(match self {
            IndexSelector::Mmap(IndexSelectorMmap { dir, is_on_disk }) => {
                NumericIndex::new_mmap(&numeric_dir(dir, field), *is_on_disk, deleted_points)?
            }
            IndexSelector::Gridstore(IndexSelectorGridstore { dir }) => {
                NumericIndex::new_gridstore(numeric_dir(dir, field), create_if_missing)?
            }
        })
    }

    fn numeric_builder<T: Encodable + Numericable + StoredValue + Send + Sync + Default, P>(
        &self,
        field: &JsonPath,
        make_mmap: fn(NumericIndexMmapBuilder<T, P>) -> FieldIndexBuilder,
        make_gridstore: fn(NumericIndexGridstoreBuilder<T, P>) -> FieldIndexBuilder,
        deleted_points: &BitSlice,
    ) -> FieldIndexBuilder
    where
        NumericIndex<T, P>: ValueIndexer<ValueType = P> + NumericIndexIntoInnerValue<T, P>,
        Vec<T>: Blob,
    {
        match self {
            IndexSelector::Mmap(IndexSelectorMmap { dir, is_on_disk }) => make_mmap(
                NumericIndex::builder_mmap(&numeric_dir(dir, field), *is_on_disk, deleted_points),
            ),
            IndexSelector::Gridstore(IndexSelectorGridstore { dir }) => {
                make_gridstore(NumericIndex::builder_gridstore(numeric_dir(dir, field)))
            }
        }
    }

    fn geo_new(
        &self,
        field: &JsonPath,
        create_if_missing: bool,
        deleted_points: &BitSlice,
    ) -> OperationResult<Option<GeoMapIndex>> {
        Ok(match self {
            IndexSelector::Mmap(IndexSelectorMmap { dir, is_on_disk }) => {
                GeoMapIndex::new_mmap(&map_dir(dir, field), *is_on_disk, deleted_points)?
            }
            IndexSelector::Gridstore(IndexSelectorGridstore { dir }) => {
                GeoMapIndex::new_gridstore(map_dir(dir, field), create_if_missing)?
            }
        })
    }

    fn geo_builder(
        &self,
        field: &JsonPath,
        make_mmap: fn(GeoMapIndexMmapBuilder) -> FieldIndexBuilder,
        make_gridstore: fn(GeoMapIndexGridstoreBuilder) -> FieldIndexBuilder,
        deleted_points: &BitSlice,
    ) -> FieldIndexBuilder {
        match self {
            IndexSelector::Mmap(IndexSelectorMmap { dir, is_on_disk }) => make_mmap(
                GeoMapIndex::builder_mmap(&map_dir(dir, field), *is_on_disk, deleted_points),
            ),
            IndexSelector::Gridstore(IndexSelectorGridstore { dir }) => {
                make_gridstore(GeoMapIndex::builder_gridstore(map_dir(dir, field)))
            }
        }
    }

    fn dir(&self) -> &Path {
        match self {
            IndexSelector::Mmap(IndexSelectorMmap { dir, is_on_disk: _ }) => dir,
            IndexSelector::Gridstore(IndexSelectorGridstore { dir }) => dir,
        }
    }

    /// Default mutability for indexes that share an on-disk format across
    /// mutability variants (e.g. bool, null): Mmap segments are immutable,
    /// Gridstore segments are appendable/mutable.
    pub fn default_mutability(&self) -> IndexMutability {
        match self {
            IndexSelector::Mmap(_) => IndexMutability::Immutable,
            IndexSelector::Gridstore(_) => IndexMutability::Mutable,
        }
    }

    pub fn null_builder(
        &self,
        field: &JsonPath,
        total_point_count: usize,
    ) -> OperationResult<FieldIndexBuilder> {
        let null_dir = null_dir(self.dir(), field);
        let builder = match self {
            IndexSelector::Mmap(_) => FieldIndexBuilder::ImmutableNullIndex(
                ImmutableNullIndex::builder(&null_dir, total_point_count)?,
            ),
            IndexSelector::Gridstore(_) => FieldIndexBuilder::MutableNullIndex(
                MutableNullIndex::builder(&null_dir, total_point_count)?,
            ),
        };
        Ok(builder)
    }

    pub fn new_null_index(
        &self,
        field: &JsonPath,
        create_if_missing: bool,
        id_tracker: &IdTrackerEnum,
        mutability: IndexMutability,
    ) -> OperationResult<Option<FieldIndex>> {
        let total_point_count = id_tracker.total_point_count();
        let null_dir = null_dir(self.dir(), field);
        // `MutableNullIndex` and `ImmutableNullIndex` share the same on-disk
        // format; stored mutability picks which in-memory wrapper to build.
        // Gridstore segments are always appendable, so the null index is
        // always mutable regardless of the stored mutability marker.
        match (self, mutability) {
            (IndexSelector::Mmap(_), IndexMutability::Immutable) => Ok(ImmutableNullIndex::open(
                &null_dir,
                total_point_count,
                id_tracker.deleted_point_bitslice(),
            )?
            .map(NullIndex::Immutable)
            .map(FieldIndex::NullIndex)),
            (IndexSelector::Mmap(_), IndexMutability::Mutable)
            | (IndexSelector::Gridstore(_), _) => {
                Ok(
                    MutableNullIndex::open(&null_dir, total_point_count, create_if_missing)?
                        .map(NullIndex::Mutable)
                        .map(FieldIndex::NullIndex),
                )
            }
        }
    }

    fn text_new(
        &self,
        field: &JsonPath,
        config: TextIndexParams,
        create_if_missing: bool,
        deleted_points: &BitSlice,
    ) -> OperationResult<Option<FullTextIndex>> {
        Ok(match self {
            IndexSelector::Mmap(IndexSelectorMmap { dir, is_on_disk }) => {
                FullTextIndex::new_mmap(text_dir(dir, field), config, *is_on_disk, deleted_points)?
            }
            IndexSelector::Gridstore(IndexSelectorGridstore { dir }) => {
                FullTextIndex::new_gridstore(text_dir(dir, field), config, create_if_missing)?
            }
        })
    }

    fn text_builder(
        &self,
        field: &JsonPath,
        config: TextIndexParams,
        deleted_points: &BitSlice,
    ) -> FieldIndexBuilder {
        match self {
            IndexSelector::Mmap(IndexSelectorMmap { dir, is_on_disk }) => {
                FieldIndexBuilder::FullTextMmapIndex(FullTextIndex::builder_mmap(
                    text_dir(dir, field),
                    config,
                    *is_on_disk,
                    deleted_points,
                ))
            }
            IndexSelector::Gridstore(IndexSelectorGridstore { dir }) => {
                FieldIndexBuilder::FullTextGridstoreIndex(FullTextIndex::builder_gridstore(
                    text_dir(dir, field),
                    config,
                ))
            }
        }
    }

    fn bool_builder(&self, field: &JsonPath) -> OperationResult<FieldIndexBuilder> {
        match self {
            IndexSelector::Mmap(IndexSelectorMmap { dir, is_on_disk: _ }) => {
                let dir = bool_dir(dir, field);
                Ok(FieldIndexBuilder::BoolMmapIndex(
                    ImmutableBoolIndex::builder(&dir)?,
                ))
            }
            // Skip Gridstore for boolean index, mmap index is simpler and is also mutable
            IndexSelector::Gridstore(IndexSelectorGridstore { dir }) => {
                let dir = bool_dir(dir, field);
                Ok(FieldIndexBuilder::BoolGridstoreIndex(
                    MutableBoolIndex::builder(&dir)?,
                ))
            }
        }
    }

    fn bool_new(
        &self,
        field: &JsonPath,
        create_if_missing: bool,
        deleted_points: &BitSlice,
        mutability: IndexMutability,
    ) -> OperationResult<Option<BoolIndex>> {
        Ok(match self {
            IndexSelector::Mmap(IndexSelectorMmap { dir, is_on_disk: _ }) => {
                let dir = bool_dir(dir, field);
                // `MutableBoolIndex` and `ImmutableBoolIndex` share the same on-disk
                // format; stored mutability picks which in-memory wrapper to build.
                match mutability {
                    IndexMutability::Immutable => {
                        ImmutableBoolIndex::open(&dir, deleted_points)?.map(BoolIndex::Immutable)
                    }
                    IndexMutability::Mutable => {
                        MutableBoolIndex::open(&dir, create_if_missing)?.map(BoolIndex::Mmap)
                    }
                }
            }
            // Skip Gridstore for boolean index, mmap index is simpler and is also mutable
            IndexSelector::Gridstore(IndexSelectorGridstore { dir }) => {
                let dir = bool_dir(dir, field);
                MutableBoolIndex::open(&dir, create_if_missing)?.map(BoolIndex::Mmap)
            }
        })
    }
}

fn map_dir(dir: &Path, field: &JsonPath) -> PathBuf {
    dir.join(format!("{}-map", &field.filename()))
}

fn numeric_dir(dir: &Path, field: &JsonPath) -> PathBuf {
    dir.join(format!("{}-numeric", &field.filename()))
}

fn text_dir(dir: &Path, field: &JsonPath) -> PathBuf {
    dir.join(format!("{}-text", &field.filename()))
}

fn bool_dir(dir: &Path, field: &JsonPath) -> PathBuf {
    dir.join(format!("{}-bool", &field.filename()))
}

fn null_dir(dir: &Path, field: &JsonPath) -> PathBuf {
    dir.join(format!("{}-null", &field.filename()))
}