engine/

engine.rs

use std::collections::HashMap;
use std::sync::Arc;

use common::{
    DakeraError, DistanceMetric, NamespaceId, PaginationCursor, QueryRequest, QueryResponse,
    Result, SearchResult,
};
use parking_lot::RwLock;
use storage::VectorStorage;

use crate::filter::evaluate_filter;
use crate::hnsw::{HnswConfig, HnswIndex};
use crate::search::brute_force_search;

/// Default vector count threshold above which HNSW is used instead of brute force
const DEFAULT_ANN_THRESHOLD: usize = 1000;

/// Convert HNSW distance back to similarity score (inverse of hnsw::similarity_to_distance)
#[inline]
fn distance_to_similarity(distance: f32, metric: DistanceMetric) -> f32 {
    match metric {
        DistanceMetric::Cosine => 1.0 - distance,
        DistanceMetric::Euclidean => -distance,
        DistanceMetric::DotProduct => -distance,
    }
}

/// Read ANN threshold from environment variable
fn ann_threshold_from_env() -> usize {
    std::env::var("DAKERA_ANN_THRESHOLD")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(DEFAULT_ANN_THRESHOLD)
}

/// Main search engine that coordinates storage and search operations
pub struct SearchEngine<S: VectorStorage + ?Sized> {
    storage: Arc<S>,
    /// Cached HNSW indices per namespace for ANN acceleration
    ann_indices: RwLock<HashMap<String, Arc<HnswIndex>>>,
    /// Vector count threshold above which HNSW is used
    ann_threshold: usize,
}

impl<S: VectorStorage + ?Sized> SearchEngine<S> {
    pub fn new(storage: Arc<S>) -> Self {
        Self {
            storage,
            ann_indices: RwLock::new(HashMap::new()),
            ann_threshold: ann_threshold_from_env(),
        }
    }

    /// Perform vector search in a namespace
    pub async fn search(
        &self,
        namespace: &NamespaceId,
        request: &QueryRequest,
    ) -> Result<QueryResponse> {
        // Check if namespace exists
        if !self.storage.namespace_exists(namespace).await? {
            return Err(DakeraError::NamespaceNotFound(namespace.clone()));
        }

        // Validate query dimension against namespace dimension
        if let Some(expected_dim) = self.storage.dimension(namespace).await? {
            if request.vector.len() != expected_dim {
                return Err(DakeraError::DimensionMismatch {
                    expected: expected_dim,
                    actual: request.vector.len(),
                });
            }
        }

        // Determine if we can use ANN acceleration:
        // - No filter (HNSW doesn't support metadata filtering)
        // - No cursor pagination (HNSW returns top-k directly)
        // - Namespace has enough vectors to justify index overhead
        let use_ann =
            request.filter.is_none() && request.cursor.is_none() && self.ann_threshold > 0;

        if use_ann {
            let count = self.storage.count(namespace).await?;
            if count > self.ann_threshold {
                return self.ann_search(namespace, request, count).await;
            }
        }

        // Fall through to brute-force search
        self.brute_force_path(namespace, request).await
    }

    /// Brute-force search path (original behavior)
    async fn brute_force_path(
        &self,
        namespace: &NamespaceId,
        request: &QueryRequest,
    ) -> Result<QueryResponse> {
        let vectors = self.storage.get_all(namespace).await?;

        let filtered_vectors: Vec<_> = if let Some(ref filter) = request.filter {
            vectors
                .into_iter()
                .filter(|v| evaluate_filter(filter, v.metadata.as_ref()))
                .collect()
        } else {
            vectors
        };

        let cursor = request
            .cursor
            .as_ref()
            .and_then(|c| PaginationCursor::decode(c));

        tracing::debug!(
            namespace = %namespace,
            vector_count = filtered_vectors.len(),
            top_k = request.top_k,
            metric = ?request.distance_metric,
            has_filter = request.filter.is_some(),
            has_cursor = cursor.is_some(),
            "Performing brute-force search"
        );

        let response = brute_force_search(
            &request.vector,
            &filtered_vectors,
            request.top_k,
            request.distance_metric,
            request.include_metadata,
            request.include_vectors,
            cursor.as_ref(),
        );

        Ok(response)
    }

    /// ANN search path using cached HNSW index
    async fn ann_search(
        &self,
        namespace: &NamespaceId,
        request: &QueryRequest,
        vector_count: usize,
    ) -> Result<QueryResponse> {
        // Get or build the HNSW index for this namespace
        let index = self
            .get_or_build_index(namespace, request.distance_metric)
            .await?;

        tracing::debug!(
            namespace = %namespace,
            vector_count = vector_count,
            top_k = request.top_k,
            metric = ?request.distance_metric,
            "Performing ANN search (HNSW)"
        );

        // Search the HNSW index — returns (VectorId, distance)
        let hnsw_results = index.search(&request.vector, request.top_k);

        // If we need metadata or vectors, fetch them from storage
        let need_fetch = request.include_metadata || request.include_vectors;
        let fetched = if need_fetch && !hnsw_results.is_empty() {
            let ids: Vec<String> = hnsw_results.iter().map(|(id, _)| id.clone()).collect();
            let vectors = self.storage.get(namespace, &ids).await?;
            let map: HashMap<String, _> = vectors.into_iter().map(|v| (v.id.clone(), v)).collect();
            Some(map)
        } else {
            None
        };

        // Convert HNSW results to SearchResults
        let results: Vec<SearchResult> = hnsw_results
            .into_iter()
            .map(|(id, distance)| {
                let score = distance_to_similarity(distance, request.distance_metric);
                let (metadata, vector) = if let Some(ref map) = fetched {
                    if let Some(v) = map.get(&id) {
                        (
                            if request.include_metadata {
                                v.metadata.clone()
                            } else {
                                None
                            },
                            if request.include_vectors {
                                Some(v.values.clone())
                            } else {
                                None
                            },
                        )
                    } else {
                        (None, None)
                    }
                } else {
                    (None, None)
                };
                SearchResult {
                    id,
                    score,
                    metadata,
                    vector,
                }
            })
            .collect();

        Ok(QueryResponse {
            results,
            next_cursor: None,
            has_more: Some(false),
            search_time_ms: 0, // caller typically overwrites this
        })
    }

    /// Get cached HNSW index or build one from storage
    async fn get_or_build_index(
        &self,
        namespace: &NamespaceId,
        metric: DistanceMetric,
    ) -> Result<Arc<HnswIndex>> {
        // Fast path: check read lock
        {
            let indices = self.ann_indices.read();
            if let Some(index) = indices.get(namespace.as_str()) {
                return Ok(Arc::clone(index));
            }
        }

        // Slow path: build index from storage
        tracing::info!(namespace = %namespace, "Building HNSW index for ANN acceleration");
        let vectors = self.storage.get_all(namespace).await?;

        let config = HnswConfig::default().with_distance_metric(metric);
        let index = HnswIndex::with_config(config);

        for v in &vectors {
            index.insert(v.id.clone(), v.values.clone());
        }

        let index = Arc::new(index);

        // Cache it
        {
            let mut indices = self.ann_indices.write();
            indices.insert(namespace.clone(), Arc::clone(&index));
        }

        tracing::info!(
            namespace = %namespace,
            vectors = vectors.len(),
            "HNSW index built and cached"
        );

        Ok(index)
    }

    /// Invalidate the cached HNSW index for a namespace (call after upsert/delete)
    pub fn invalidate_ann_index(&self, namespace: &NamespaceId) {
        let mut indices = self.ann_indices.write();
        if indices.remove(namespace.as_str()).is_some() {
            tracing::debug!(namespace = %namespace, "HNSW index invalidated");
        }
    }

    /// Get reference to storage
    pub fn storage(&self) -> &Arc<S> {
        &self.storage
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use common::{DistanceMetric, FilterCondition, FilterExpression, FilterValue, Vector};
    use std::collections::HashMap;
    use storage::InMemoryStorage;

    async fn setup_engine() -> (SearchEngine<InMemoryStorage>, String) {
        let storage = Arc::new(InMemoryStorage::new());
        let engine = SearchEngine::new(storage.clone());
        let namespace = "test".to_string();

        storage.ensure_namespace(&namespace).await.unwrap();
        storage
            .upsert(
                &namespace,
                vec![
                    Vector {
                        id: "v1".to_string(),
                        values: vec![1.0, 0.0, 0.0],
                        metadata: None,
                        ttl_seconds: None,
                        expires_at: None,
                    },
                    Vector {
                        id: "v2".to_string(),
                        values: vec![0.0, 1.0, 0.0],
                        metadata: None,
                        ttl_seconds: None,
                        expires_at: None,
                    },
                    Vector {
                        id: "v3".to_string(),
                        values: vec![0.707, 0.707, 0.0],
                        metadata: None,
                        ttl_seconds: None,
                        expires_at: None,
                    },
                ],
            )
            .await
            .unwrap();

        (engine, namespace)
    }

    #[tokio::test]
    async fn test_search_basic() {
        let (engine, namespace) = setup_engine().await;

        let request = QueryRequest {
            vector: vec![1.0, 0.0, 0.0],
            top_k: 2,
            distance_metric: DistanceMetric::Cosine,
            include_metadata: true,
            include_vectors: false,
            filter: None,
            cursor: None,
            consistency: Default::default(),
            staleness_config: None,
        };

        let response = engine.search(&namespace, &request).await.unwrap();

        assert_eq!(response.results.len(), 2);
        assert_eq!(response.results[0].id, "v1"); // Perfect match
    }

    #[tokio::test]
    async fn test_search_namespace_not_found() {
        let storage = Arc::new(InMemoryStorage::new());
        let engine = SearchEngine::new(storage);

        let request = QueryRequest {
            vector: vec![1.0, 0.0, 0.0],
            top_k: 5,
            distance_metric: DistanceMetric::Cosine,
            include_metadata: true,
            include_vectors: false,
            filter: None,
            cursor: None,
            consistency: Default::default(),
            staleness_config: None,
        };

        let result = engine.search(&"nonexistent".to_string(), &request).await;

        assert!(matches!(result, Err(DakeraError::NamespaceNotFound(_))));
    }

    #[tokio::test]
    async fn test_search_dimension_mismatch() {
        let (engine, namespace) = setup_engine().await;

        let request = QueryRequest {
            vector: vec![1.0, 0.0], // Wrong dimension (2 instead of 3)
            top_k: 5,
            distance_metric: DistanceMetric::Cosine,
            include_metadata: true,
            include_vectors: false,
            filter: None,
            cursor: None,
            consistency: Default::default(),
            staleness_config: None,
        };

        let result = engine.search(&namespace, &request).await;

        assert!(matches!(
            result,
            Err(DakeraError::DimensionMismatch {
                expected: 3,
                actual: 2
            })
        ));
    }

    #[tokio::test]
    async fn test_search_empty_namespace() {
        let storage = Arc::new(InMemoryStorage::new());
        let engine = SearchEngine::new(storage.clone());
        let namespace = "empty".to_string();

        storage.ensure_namespace(&namespace).await.unwrap();

        let request = QueryRequest {
            vector: vec![1.0, 0.0, 0.0],
            top_k: 5,
            distance_metric: DistanceMetric::Cosine,
            include_metadata: true,
            include_vectors: false,
            filter: None,
            cursor: None,
            consistency: Default::default(),
            staleness_config: None,
        };

        let response = engine.search(&namespace, &request).await.unwrap();

        assert!(response.results.is_empty());
    }

    #[tokio::test]
    async fn test_search_with_filter() {
        let storage = Arc::new(InMemoryStorage::new());
        let engine = SearchEngine::new(storage.clone());
        let namespace = "test".to_string();

        storage.ensure_namespace(&namespace).await.unwrap();
        storage
            .upsert(
                &namespace,
                vec![
                    Vector {
                        id: "v1".to_string(),
                        values: vec![1.0, 0.0, 0.0],
                        metadata: Some(
                            serde_json::json!({"category": "electronics", "price": 100}),
                        ),
                        ttl_seconds: None,
                        expires_at: None,
                    },
                    Vector {
                        id: "v2".to_string(),
                        values: vec![0.9, 0.1, 0.0],
                        metadata: Some(serde_json::json!({"category": "books", "price": 20})),
                        ttl_seconds: None,
                        expires_at: None,
                    },
                    Vector {
                        id: "v3".to_string(),
                        values: vec![0.8, 0.2, 0.0],
                        metadata: Some(serde_json::json!({"category": "electronics", "price": 50})),
                        ttl_seconds: None,
                        expires_at: None,
                    },
                ],
            )
            .await
            .unwrap();

        // Filter for electronics only
        let mut field = HashMap::new();
        field.insert(
            "category".to_string(),
            FilterCondition::Eq(FilterValue::String("electronics".to_string())),
        );

        let request = QueryRequest {
            vector: vec![1.0, 0.0, 0.0],
            top_k: 10,
            distance_metric: DistanceMetric::Cosine,
            include_metadata: true,
            include_vectors: false,
            filter: Some(FilterExpression::Field { field }),
            cursor: None,
            consistency: Default::default(),
            staleness_config: None,
        };

        let response = engine.search(&namespace, &request).await.unwrap();

        // Should only return v1 and v3 (electronics)
        assert_eq!(response.results.len(), 2);
        assert!(response
            .results
            .iter()
            .all(|r| r.id == "v1" || r.id == "v3"));
    }

    #[tokio::test]
    async fn test_search_with_numeric_filter() {
        let storage = Arc::new(InMemoryStorage::new());
        let engine = SearchEngine::new(storage.clone());
        let namespace = "test".to_string();

        storage.ensure_namespace(&namespace).await.unwrap();
        storage
            .upsert(
                &namespace,
                vec![
                    Vector {
                        id: "v1".to_string(),
                        values: vec![1.0, 0.0, 0.0],
                        metadata: Some(serde_json::json!({"price": 100})),
                        ttl_seconds: None,
                        expires_at: None,
                    },
                    Vector {
                        id: "v2".to_string(),
                        values: vec![0.9, 0.1, 0.0],
                        metadata: Some(serde_json::json!({"price": 20})),
                        ttl_seconds: None,
                        expires_at: None,
                    },
                    Vector {
                        id: "v3".to_string(),
                        values: vec![0.8, 0.2, 0.0],
                        metadata: Some(serde_json::json!({"price": 50})),
                        ttl_seconds: None,
                        expires_at: None,
                    },
                ],
            )
            .await
            .unwrap();

        // Filter for price < 60
        let mut field = HashMap::new();
        field.insert(
            "price".to_string(),
            FilterCondition::Lt(FilterValue::Number(60.0)),
        );

        let request = QueryRequest {
            vector: vec![1.0, 0.0, 0.0],
            top_k: 10,
            distance_metric: DistanceMetric::Cosine,
            include_metadata: true,
            include_vectors: false,
            filter: Some(FilterExpression::Field { field }),
            cursor: None,
            consistency: Default::default(),
            staleness_config: None,
        };

        let response = engine.search(&namespace, &request).await.unwrap();

        // Should only return v2 (20) and v3 (50)
        assert_eq!(response.results.len(), 2);
        assert!(response
            .results
            .iter()
            .all(|r| r.id == "v2" || r.id == "v3"));
    }
}