qdrant-edge 0.7.2

use std::collections::HashMap;
use std::hash::{DefaultHasher, Hash, Hasher};

use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use smallvec::SmallVec;

use crate::segment::data_types::named_vectors::NamedVectors;
use crate::segment::data_types::vectors::VectorInternal;
use crate::segment::index::field_index::FieldIndex;
use crate::segment::types::{PayloadKeyType, VectorNameBuf};

pub type IndexesMap = HashMap<PayloadKeyType, Vec<FieldIndex>>;

/// A container for JSON values, optimized for the common case of a single value.
pub type MultiValue<T> = SmallVec<[T; 1]>;

pub fn check_is_empty<'a>(values: impl IntoIterator<Item = &'a Value>) -> bool {
    values.into_iter().all(|x| match x {
        Value::Null => true,
        Value::Array(arr) => arr.is_empty(),
        Value::Bool(_) | Value::Number(_) | Value::String(_) | Value::Object(_) => false,
    })
}

pub fn check_is_null<'a>(values: impl IntoIterator<Item = &'a Value>) -> bool {
    values.into_iter().any(|x| x.is_null())
    // { "a": [ { "b": null }, { "b": 1 } ] } => true
    // { "a": [ { "b": 1 }, { "b": null } ] } => true
    // { "a": [ { "b": 1 }, { "b": 2 } ] } => false
}

pub fn rev_range(a: usize, b: usize) -> impl Iterator<Item = usize> {
    (b + 1..=a).rev()
}

// Merge source map into destination map
pub fn merge_map(
    dest: &mut serde_json::Map<String, Value>,
    source: &serde_json::Map<String, Value>,
) {
    for (key, value) in source {
        match value {
            Value::Null => dest.remove(key),
            Value::Bool(_)
            | Value::Number(_)
            | Value::String(_)
            | Value::Array(_)
            | Value::Object(_) => dest.insert(key.to_owned(), value.to_owned()),
        };
    }
}

pub fn transpose_map_into_named_vector<TVector: Into<VectorInternal>>(
    map: HashMap<VectorNameBuf, Vec<TVector>>,
) -> Vec<NamedVectors<'static>> {
    let mut result = Vec::new();
    for (key, values) in map {
        result.resize_with(values.len(), NamedVectors::default);
        for (i, value) in values.into_iter().enumerate() {
            result[i].insert(key.clone(), value.into());
        }
    }
    result
}

/// Hashes an iterator of unique items in an order-independent way.
///
/// The order of items does not affect the resulting hash.
/// Assumes that each item is unique. Suitable for hashing, e.g., sets and maps,
/// but not for vectors or multisets.
#[inline(always)]
pub fn unordered_hash_unique<T: Hash, I: Iterator<Item = T>, H: Hasher>(state: &mut H, iter: I) {
    iter.fold(0u64, |res, item| {
        let mut hasher = DefaultHasher::new();
        item.hash(&mut hasher);
        // NOTE: xor is used, duplicates will cancel each other out
        res ^ hasher.finish()
    })
    .hash(state);
}

/// Deserializer helper for `Option<Vec<T>>` that allows deserializing both single and an array of values.
///
/// Use via `#[serde(with = "MaybeOneOrMany")]` and `#[schemars(with="MaybeOneOrMany<T>")]` field attributes
pub struct MaybeOneOrMany<T>(pub Option<Vec<T>>);

impl<T: Serialize> MaybeOneOrMany<T> {
    pub fn serialize<S>(value: &Option<Vec<T>>, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        value.serialize(serializer)
    }
}

impl<'de, T: Deserialize<'de>> MaybeOneOrMany<T> {
    pub fn deserialize<D>(deserializer: D) -> Result<Option<Vec<T>>, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        use serde_untagged::UntaggedEnumVisitor;

        UntaggedEnumVisitor::new()
            .unit(|| Ok(None))
            .seq(|x| x.deserialize().map(Some))
            .map(|x| x.deserialize().map(|x| vec![x]).map(Some))
            .deserialize(deserializer)
    }
}

impl<T: JsonSchema> JsonSchema for MaybeOneOrMany<T> {
    fn json_schema(generator: &mut schemars::r#gen::SchemaGenerator) -> schemars::schema::Schema {
        use schemars::schema::SchemaObject;

        #[derive(JsonSchema)]
        #[serde(untagged)]
        enum OneOrMany<T> {
            _One(T),
            _Many(Vec<T>),
            _None(()),
        }

        let schema: SchemaObject = <OneOrMany<T>>::json_schema(generator).into();
        schema.into()
    }

    fn schema_name() -> String {
        <Vec<T>>::schema_name()
    }

    fn is_referenceable() -> bool {
        false
    }
}

#[cfg(test)]
mod tests {
    use schemars::{JsonSchema, schema_for};
    use serde::{Deserialize, Serialize};

    use crate::segment::common::utils::MaybeOneOrMany;

    #[test]
    fn test_deserialize_one_or_many() {
        #[derive(Serialize, Deserialize)]
        struct Test {
            #[serde(with = "MaybeOneOrMany")]
            data: Option<Vec<Inner>>,
        }

        #[derive(Serialize, Deserialize)]
        struct Inner {
            key: String,
        }

        let res = serde_json::from_str::<Test>(
            r#"
            {
                "data": null
            }
            "#,
        )
        .unwrap();

        assert!(res.data.is_none());

        let res = serde_json::from_str::<Test>(
            r#"
            {
                "data": {
                    "key": "value"
                }
            }
            "#,
        )
        .unwrap();

        assert_eq!(res.data.as_ref().unwrap().len(), 1);
        assert_eq!(res.data.as_ref().unwrap()[0].key, "value".to_string());

        let res = serde_json::from_str::<Test>(
            r#"
            {
                "data": [
                    {
                        "key": "value"
                    }
                ]
            }
            "#,
        )
        .unwrap();

        assert_eq!(res.data.as_ref().unwrap().len(), 1);
        assert_eq!(res.data.as_ref().unwrap()[0].key, "value".to_string());
    }

    #[test]
    fn test_schema_one_or_many() {
        #[derive(JsonSchema)]
        struct Test {
            #[schemars(with = "MaybeOneOrMany<String>")]
            _field: Option<Vec<String>>,
        }

        let mut field_schema = dbg!(
            schemars::schema_for!(Test)
                .schema
                .object
                .unwrap()
                .properties
                .remove("_field")
                .unwrap()
                .into_object(),
        );

        assert!(field_schema.subschemas.is_some());

        let any_of = field_schema.subschemas().any_of.clone().unwrap();

        assert_eq!(any_of.len(), 3);
        assert_eq!(
            any_of[0].clone().into_object().instance_type,
            schema_for!(String).schema.instance_type
        );
        assert_eq!(
            any_of[1].clone().into_object().array,
            schema_for!(Vec<String>).schema.array
        );
        assert_eq!(
            any_of[2].clone().into_object().instance_type,
            schema_for!(()).schema.instance_type
        );
    }
}