//! Unified vector search engine
//!
//! Provides a `VectorSearchEngine` that abstracts over different vector
//! search strategies:
//!
//! - **HNSW**: In-memory approximate nearest neighbor search (existing)
//! - **Lance**: Lance native scanner-based ANN search (new)
//! - **Hybrid**: HNSW for search + Lance for persistence (recommended)
//!
//! # Architecture
//!
//! ```text
//!              ┌───────────────────────┐
//!              │  VectorSearchEngine   │
//!              └──────────┬────────────┘
//!                ┌────────┼────────┐
//!                │        │        │
//!        ┌───────▼───┐ ┌──▼──┐ ┌───▼────┐
//!        │   HNSW    │ │Lance│ │ Hybrid  │
//!        │ (in-mem)  │ │(ANN)│ │HNSW+L. │
//!        └───────────┘ └─────┘ └────────┘
//! ```

use std::sync::Arc;

use crate::storage::backend::{StorageError, StorageResult};
use crate::storage::vector_backend::{VectorBackend, VectorSearchResult, VectorIndexConfig};
use crate::vector::hnsw::HNSWIndex;

/// Vector search strategy configuration.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum VectorSearchStrategy {
    /// Use HNSW in-memory index only (existing behavior).
    /// Vectors are stored in HNSW nodes, persisted via bincode.
    Hnsw,

    /// Use Lance native scanner-based search only.
    /// Vectors are stored in Lance dataset's embedding column.
    Lance,

    /// Use HNSW for fast in-memory search + Lance for durable persistence.
    /// Vectors are stored in both HNSW (for search) and Lance (for persistence).
    /// On startup, HNSW is rebuilt from the Lance dataset.
    Hybrid,
}

impl Default for VectorSearchStrategy {
    fn default() -> Self {
        Self::Hnsw // default to existing behavior
    }
}

/// Unified vector search engine.
///
/// Wraps either HNSW, a `VectorBackend` (Lance), or both,
/// and provides a single search interface.
pub struct VectorSearchEngine {
    /// HNSW index (always available for Hnsw and Hybrid modes)
    hnsw: Option<Arc<HNSWIndex>>,

    /// Vector backend (available for Lance and Hybrid modes)
    vector_backend: Option<Arc<dyn VectorBackend>>,

    /// Active strategy
    strategy: VectorSearchStrategy,
}

impl VectorSearchEngine {
    /// Create an HNSW-only engine (existing behavior).
    pub fn hnsw(index: Arc<HNSWIndex>) -> Self {
        Self {
            hnsw: Some(index),
            vector_backend: None,
            strategy: VectorSearchStrategy::Hnsw,
        }
    }

    /// Create a Lance-only engine.
    pub fn lance(backend: Arc<dyn VectorBackend>) -> Self {
        Self {
            hnsw: None,
            vector_backend: Some(backend),
            strategy: VectorSearchStrategy::Lance,
        }
    }

    /// Create a hybrid engine: HNSW for search + Lance for persistence.
    pub fn hybrid(index: Arc<HNSWIndex>, backend: Arc<dyn VectorBackend>) -> Self {
        Self {
            hnsw: Some(index),
            vector_backend: Some(backend),
            strategy: VectorSearchStrategy::Hybrid,
        }
    }

    /// Get the active search strategy.
    pub fn strategy(&self) -> &VectorSearchStrategy {
        &self.strategy
    }

    /// Get a reference to the HNSW index (if available).
    pub fn hnsw_index(&self) -> Option<&Arc<HNSWIndex>> {
        self.hnsw.as_ref()
    }

    /// Get a reference to the vector backend (if available).
    pub fn vector_backend(&self) -> Option<&Arc<dyn VectorBackend>> {
        self.vector_backend.as_ref()
    }

    /// Insert a vector into the search engine.
    ///
    /// For HNSW-only: inserts into HNSW.
    /// For Lance-only: inserts into the vector backend.
    /// For Hybrid: inserts into both.
    pub fn insert(&self, id: &str, vector: &[f32]) -> StorageResult<()> {
        match self.strategy {
            VectorSearchStrategy::Hnsw => {
                let hnsw = self.hnsw.as_ref()
                    .ok_or_else(|| StorageError::Backend("HNSW index not available".into()))?;
                hnsw.insert(id.to_string(), vector.to_vec())
                    .map_err(|e| StorageError::Backend(format!("HNSW insert error: {}", e)))?;
            }
            VectorSearchStrategy::Lance => {
                let backend = self.vector_backend.as_ref()
                    .ok_or_else(|| StorageError::Backend("Vector backend not available".into()))?;
                backend.store_vector(id, vector)?;
            }
            VectorSearchStrategy::Hybrid => {
                // Insert into HNSW for search
                if let Some(ref hnsw) = self.hnsw {
                    hnsw.insert(id.to_string(), vector.to_vec())
                        .map_err(|e| StorageError::Backend(format!("HNSW insert error: {}", e)))?;
                }
                // Insert into Lance for persistence
                if let Some(ref backend) = self.vector_backend {
                    backend.store_vector(id, vector)?;
                }
            }
        }
        Ok(())
    }

    /// Search for the k nearest neighbors.
    ///
    /// For HNSW/Hybrid: uses HNSW (fast in-memory).
    /// For Lance: uses the vector backend (scanner-based ANN).
    pub fn search(&self, query: &[f32], k: usize) -> StorageResult<Vec<VectorSearchResult>> {
        match self.strategy {
            VectorSearchStrategy::Hnsw | VectorSearchStrategy::Hybrid => {
                let hnsw = self.hnsw.as_ref()
                    .ok_or_else(|| StorageError::Backend("HNSW index not available".into()))?;
                let results = hnsw.search_with_ids(query, k)
                    .map_err(|e| StorageError::Backend(format!("HNSW search error: {}", e)))?;
                Ok(results
                    .into_iter()
                    .map(|(id, dist)| VectorSearchResult::new(id, dist))
                    .collect())
            }
            VectorSearchStrategy::Lance => {
                let backend = self.vector_backend.as_ref()
                    .ok_or_else(|| StorageError::Backend("Vector backend not available".into()))?;
                backend.search_nearest(query, k)
            }
        }
    }

    /// Search with ID filter.
    pub fn search_filtered(
        &self,
        query: &[f32],
        k: usize,
        allowed_ids: &[String],
    ) -> StorageResult<Vec<VectorSearchResult>> {
        match self.strategy {
            VectorSearchStrategy::Hnsw | VectorSearchStrategy::Hybrid => {
                // HNSW doesn't natively support filtering, search more then filter
                let hnsw = self.hnsw.as_ref()
                    .ok_or_else(|| StorageError::Backend("HNSW index not available".into()))?;
                let results = hnsw.search_with_ids(query, k * 4)
                    .map_err(|e| StorageError::Backend(format!("HNSW search error: {}", e)))?;
                let mut filtered: Vec<VectorSearchResult> = results
                    .into_iter()
                    .filter(|(id, _)| allowed_ids.contains(id))
                    .map(|(id, dist)| VectorSearchResult::new(id, dist))
                    .collect();
                filtered.truncate(k);
                Ok(filtered)
            }
            VectorSearchStrategy::Lance => {
                let backend = self.vector_backend.as_ref()
                    .ok_or_else(|| StorageError::Backend("Vector backend not available".into()))?;
                backend.search_nearest_filtered(query, k, allowed_ids)
            }
        }
    }

    /// Filtered ANN search using a `FilterExpr` AST.
    ///
    /// For Lance backend: uses native Lance push-down filtering combined with
    /// vector nearest-neighbor search for optimal performance.
    /// For HNSW: filters candidates after the search.
    pub fn search_with_expr(
        &self,
        query: &[f32],
        k: usize,
        expr: &crate::storage::filter_expr::FilterExpr,
        doc_store: Option<&dyn crate::storage::backend::StorageBackend>,
    ) -> StorageResult<Vec<VectorSearchResult>> {
        match self.strategy {
            VectorSearchStrategy::Lance => {
                let backend = self.vector_backend.as_ref()
                    .ok_or_else(|| StorageError::Backend("Vector backend not available".into()))?;
                backend.search_nearest_with_expr(query, k, expr, doc_store)
            }
            VectorSearchStrategy::Hnsw | VectorSearchStrategy::Hybrid => {
                // For HNSW: pre-filter using doc_store, then search with allowed IDs
                if let Some(store) = doc_store {
                    let filtered_nodes = store.query_expr(
                        expr,
                        None,
                        crate::storage::backend::SortOrder::Asc,
                        0, 0,
                    )?;
                    let allowed_ids: Vec<String> = filtered_nodes.iter().map(|n| n.id.clone()).collect();
                    if allowed_ids.is_empty() {
                        return Ok(Vec::new());
                    }
                    self.search_filtered(query, k, &allowed_ids)
                } else {
                    // No doc store → unfiltered search
                    self.search(query, k)
                }
            }
        }
    }

    /// Delete a vector from the search engine.
    pub fn delete(&self, id: &str) -> StorageResult<()> {
        match self.strategy {
            VectorSearchStrategy::Hnsw => {
                if let Some(ref hnsw) = self.hnsw {
                    let _ = hnsw.delete(id);
                }
            }
            VectorSearchStrategy::Lance => {
                if let Some(ref backend) = self.vector_backend {
                    backend.delete_vector(id)?;
                }
            }
            VectorSearchStrategy::Hybrid => {
                if let Some(ref hnsw) = self.hnsw {
                    let _ = hnsw.delete(id);
                }
                if let Some(ref backend) = self.vector_backend {
                    backend.delete_vector(id)?;
                }
            }
        }
        Ok(())
    }

    /// Get a vector by ID.
    pub fn get_vector(&self, id: &str) -> StorageResult<Option<Vec<f32>>> {
        match self.strategy {
            VectorSearchStrategy::Hnsw | VectorSearchStrategy::Hybrid => {
                if let Some(ref hnsw) = self.hnsw {
                    if let Some(vec) = hnsw.get(id) {
                        return Ok(Some(vec));
                    }
                }
                // Fallback to Lance backend if available
                if let Some(ref backend) = self.vector_backend {
                    return backend.get_vector(id);
                }
                Ok(None)
            }
            VectorSearchStrategy::Lance => {
                let backend = self.vector_backend.as_ref()
                    .ok_or_else(|| StorageError::Backend("Vector backend not available".into()))?;
                backend.get_vector(id)
            }
        }
    }

    /// Get number of vectors.
    pub fn vector_count(&self) -> usize {
        match self.strategy {
            VectorSearchStrategy::Hnsw | VectorSearchStrategy::Hybrid => {
                self.hnsw.as_ref().map(|h| h.len()).unwrap_or(0)
            }
            VectorSearchStrategy::Lance => {
                self.vector_backend.as_ref()
                    .and_then(|b| b.vector_count().ok())
                    .unwrap_or(0)
            }
        }
    }

    /// Flush the search engine state.
    pub fn flush(&self) -> StorageResult<()> {
        if let Some(ref backend) = self.vector_backend {
            backend.flush()?;
        }
        Ok(())
    }

    /// Rebuild HNSW index from the vector backend.
    ///
    /// Useful for Hybrid mode on startup: load vectors from Lance
    /// and insert them into HNSW for fast search.
    pub fn rebuild_hnsw_from_backend(&self) -> StorageResult<usize> {
        let hnsw = self.hnsw.as_ref()
            .ok_or_else(|| StorageError::Backend("HNSW index not available".into()))?;
        let backend = self.vector_backend.as_ref()
            .ok_or_else(|| StorageError::Backend("Vector backend not available".into()))?;

        let ids = backend.vector_ids()?;
        let mut count = 0;

        for id in &ids {
            if let Some(vector) = backend.get_vector(id)? {
                if hnsw.insert(id.clone(), vector).is_ok() {
                    count += 1;
                }
            }
        }

        Ok(count)
    }

    /// Create a vector index on the Lance backend for accelerated search.
    pub fn create_lance_index(&self, config: &VectorIndexConfig) -> StorageResult<()> {
        let backend = self.vector_backend.as_ref()
            .ok_or_else(|| StorageError::Backend("Vector backend not available".into()))?;
        backend.create_index(config)
    }

    /// Get the backend name.
    pub fn backend_name(&self) -> &str {
        match self.strategy {
            VectorSearchStrategy::Hnsw => "hnsw",
            VectorSearchStrategy::Lance => {
                self.vector_backend.as_ref()
                    .map(|b| b.backend_name())
                    .unwrap_or("lance")
            }
            VectorSearchStrategy::Hybrid => "hybrid(hnsw+lance)",
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::HNSWConfig;

    #[test]
    fn test_hnsw_only_engine() {
        let config = HNSWConfig::new(4).with_max_elements(100);
        let hnsw = Arc::new(HNSWIndex::new(config));

        let engine = VectorSearchEngine::hnsw(hnsw);
        assert_eq!(engine.backend_name(), "hnsw");

        // Insert vectors
        engine.insert("v1", &[1.0, 0.0, 0.0, 0.0]).unwrap();
        engine.insert("v2", &[0.0, 1.0, 0.0, 0.0]).unwrap();
        engine.insert("v3", &[0.9, 0.1, 0.0, 0.0]).unwrap();

        assert_eq!(engine.vector_count(), 3);

        // Search
        let results = engine.search(&[1.0, 0.0, 0.0, 0.0], 2).unwrap();
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].id, "v1"); // exact match first
    }

    #[test]
    fn test_hybrid_engine_with_mmap() {
        use crate::storage::vector_backend::MmapVectorBackend;
        use crate::types::StorageConfig;
        use tempfile::tempdir;

        let dir = tempdir().unwrap();
        let mmap_config = StorageConfig {
            path: dir.path().to_string_lossy().to_string(),
            ..Default::default()
        };
        let mmap_backend = Arc::new(MmapVectorBackend::open_or_create(mmap_config, 3).unwrap());

        let hnsw_config = HNSWConfig::new(3).with_max_elements(100);
        let hnsw = Arc::new(HNSWIndex::new(hnsw_config));

        let engine = VectorSearchEngine::hybrid(hnsw, mmap_backend.clone());
        assert_eq!(engine.backend_name(), "hybrid(hnsw+lance)");

        // Insert vectors — should go to both HNSW and mmap
        engine.insert("a", &[1.0, 0.0, 0.0]).unwrap();
        engine.insert("b", &[0.0, 1.0, 0.0]).unwrap();

        assert_eq!(engine.vector_count(), 2);

        // Search uses HNSW
        let results = engine.search(&[1.0, 0.0, 0.0], 1).unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, "a");

        // Get vector — available from both
        let vec = engine.get_vector("a").unwrap();
        assert!(vec.is_some());
        assert_eq!(vec.unwrap(), vec![1.0, 0.0, 0.0]);
    }

    #[test]
    fn test_hnsw_rebuild_from_backend() {
        use crate::storage::vector_backend::MmapVectorBackend;
        use crate::types::StorageConfig;
        use tempfile::tempdir;

        let dir = tempdir().unwrap();
        let mmap_config = StorageConfig {
            path: dir.path().to_string_lossy().to_string(),
            ..Default::default()
        };
        let mmap_backend = Arc::new(MmapVectorBackend::open_or_create(mmap_config, 3).unwrap());

        // Store vectors in mmap (simulating Lance persistence)
        mmap_backend.store_vector("x", &[1.0, 0.0, 0.0]).unwrap();
        mmap_backend.store_vector("y", &[0.0, 1.0, 0.0]).unwrap();
        mmap_backend.store_vector("z", &[0.0, 0.0, 1.0]).unwrap();

        // Create fresh HNSW (empty)
        let hnsw_config = HNSWConfig::new(3).with_max_elements(100);
        let hnsw = Arc::new(HNSWIndex::new(hnsw_config));
        assert_eq!(hnsw.len(), 0);

        let engine = VectorSearchEngine::hybrid(hnsw, mmap_backend);

        // Rebuild HNSW from the vector backend
        let count = engine.rebuild_hnsw_from_backend().unwrap();
        assert_eq!(count, 3);
        assert_eq!(engine.vector_count(), 3);

        // Search should work now
        let results = engine.search(&[1.0, 0.0, 0.0], 1).unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, "x");
    }
}
