//! `USearch`-backed HNSW vector index implementation.
//!
//! [`UsearchIndex`] wraps the `usearch` crate's [`usearch::Index`] behind
//! the [`VectorIndex`] trait, providing the production-grade Approximate
//! Nearest Neighbor search backend for tidalDB.
//!
//! # Design decisions
//!
//! - **f32 input, internal quantization**: All vectors are accepted as `f32`
//!   slices. `USearch` handles quantization internally (to F16 or I8) based on
//!   the `ScalarKind` specified at construction time.
//! - **`total_slots` tracking**: `USearch`'s `size()` excludes tombstoned entries
//!   after `remove()`. We maintain an `AtomicUsize` counter to provide the total
//!   occupied graph slot count (including tombstones) via `len()`.
//! - **No per-query `ef_search` override**: `USearch`'s `change_expansion_search()`
//!   mutates global index state (not per-query). We accept the parameter for
//!   trait compliance but use the index-level default set at construction time.
//!   See TODO below.
//! - **Thread safety**: `usearch::Index` implements `Send + Sync` (backed by
//!   C++ thread-safe HNSW). `AtomicUsize` is trivially `Send + Sync`. Therefore
//!   `UsearchIndex` is `Send + Sync` without any unsafe impl blocks.

use std::path::Path;
use std::sync::atomic::{AtomicUsize, Ordering};

use super::{DistanceMetric, QuantizationLevel, VectorSearchResult};
use super::{VectorError, VectorId, VectorIndex, VectorIndexConfig};

/// Production HNSW vector index backed by `USearch`.
///
/// Wraps [`usearch::Index`] behind the [`VectorIndex`] trait so the rest
/// of tidalDB never imports `USearch` directly. If we swap to a different
/// HNSW implementation, only this module changes.
///
/// # Tombstone behavior
///
/// `USearch`'s `remove()` uses lazy tombstoning internally: the graph node
/// is marked as removed and excluded from search results, but the slot
/// is not immediately reclaimed. However, `USearch`'s `size()` already
/// decrements on remove, reflecting only live entries. We track the total
/// number of entries ever added (minus reuses) in `total_slots` so that
/// `len()` can report the true graph size (including tombstoned slots),
/// while `len_live()` reports `size()` (only live entries). This enables
/// accurate `tombstone_ratio()` for triggering compaction decisions.
pub struct UsearchIndex {
    inner: usearch::Index,
    config: VectorIndexConfig,
    /// Tracks the total number of occupied slots in the HNSW graph,
    /// including tombstoned entries. `USearch`'s `size()` excludes removed
    /// entries, so we maintain this separately.
    ///
    /// Incremented on `insert()` for new keys, never decremented on
    /// `delete()` (because the graph slot remains allocated).
    ///
    /// Ordering rationale:
    /// - `Relaxed` is sufficient because this counter is a best-effort
    ///   statistic. It does not gate correctness of search results or
    ///   persistence. A ranking query that reads a slightly stale count
    ///   suffers no correctness issue.
    total_slots: AtomicUsize,
}

impl UsearchIndex {
    /// Create a new, empty `USearch` HNSW index with the given configuration.
    ///
    /// # Errors
    ///
    /// Returns [`VectorError::Backend`] if `USearch` fails to allocate the index
    /// (e.g., unsupported metric/quantization combination).
    pub fn new(config: VectorIndexConfig) -> Result<Self, VectorError> {
        let metric = match config.metric {
            DistanceMetric::L2 => usearch::MetricKind::L2sq,
            DistanceMetric::InnerProduct => usearch::MetricKind::IP,
        };
        let quantization = match config.quantization {
            QuantizationLevel::F32 => usearch::ScalarKind::F32,
            QuantizationLevel::F16 => usearch::ScalarKind::F16,
            QuantizationLevel::Int8 => usearch::ScalarKind::I8,
        };

        let options = usearch::IndexOptions {
            dimensions: config.dimensions,
            metric,
            quantization,
            connectivity: config.connectivity,
            expansion_add: config.ef_construction,
            expansion_search: config.ef_search,
            multi: false,
        };

        let inner = usearch::new_index(&options)
            .map_err(|e| VectorError::Backend(format!("USearch init failed: {e}")))?;

        Ok(Self {
            inner,
            config,
            total_slots: AtomicUsize::new(0),
        })
    }

    /// Validate that a vector's dimensionality matches the index configuration.
    const fn validate_dimensions(&self, vec: &[f32]) -> Result<(), VectorError> {
        if vec.len() != self.config.dimensions {
            return Err(VectorError::DimensionMismatch {
                expected: self.config.dimensions,
                got: vec.len(),
            });
        }
        Ok(())
    }

    /// Convert a `&Path` to a `&str`, returning a `VectorError::Io` on non-UTF-8 paths.
    fn path_to_str(path: &Path) -> Result<&str, VectorError> {
        path.to_str().ok_or_else(|| {
            VectorError::Io(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "non-UTF-8 path",
            ))
        })
    }
}

impl VectorIndex for UsearchIndex {
    fn insert(&self, id: VectorId, embedding: &[f32]) -> Result<(), VectorError> {
        self.validate_dimensions(embedding)?;

        // Track whether this is a new key or a replacement.
        // New keys allocate a new graph slot; replacements reuse the existing slot.
        let is_new = !self.inner.contains(id);

        self.inner
            .add(id, embedding)
            .map_err(|e| VectorError::Backend(format!("USearch insert failed: {e}")))?;

        if is_new {
            self.total_slots.fetch_add(1, Ordering::Relaxed);
        }

        Ok(())
    }

    /// Search for the `k` nearest neighbors using HNSW approximate search.
    ///
    /// The `ef_search` parameter is accepted for trait compliance. If it
    /// differs from the index-level default, a debug log is emitted, but
    /// the index-level default is used.
    ///
    /// # Per-query `ef_search` override (TODO)
    ///
    /// `USearch`'s `change_expansion_search()` mutates global state and is
    /// not safe for concurrent queries with different `ef_search` values.
    /// A proper solution would require either:
    ///
    /// 1. A mutex around search + expansion change (unacceptable on the hot path), or
    /// 2. `USearch` adding a per-query expansion parameter to their search API.
    ///
    /// For now, we use the index-level default set at construction time.
    fn search(
        &self,
        query: &[f32],
        k: usize,
        ef_search: usize,
    ) -> Result<Vec<VectorSearchResult>, VectorError> {
        self.validate_dimensions(query)?;

        if ef_search != self.config.ef_search {
            tracing::warn!(
                requested = ef_search,
                default = self.config.ef_search,
                "per-query ef_search override not supported; using index-level default (may affect recall)"
            );
        }

        let matches = self
            .inner
            .search(query, k)
            .map_err(|e| VectorError::Backend(format!("USearch search failed: {e}")))?;

        Ok(matches
            .keys
            .into_iter()
            .zip(matches.distances)
            .map(|(id, distance)| VectorSearchResult { id, distance })
            .collect())
    }

    /// Filtered search: only consider vectors where `filter(id)` returns `true`.
    ///
    /// `USearch`'s `filtered_search` uses a trampoline to pass a Rust closure
    /// through the C++ layer. The predicate is evaluated during graph traversal,
    /// so non-matching vectors are skipped without computing their distance.
    ///
    /// See `search()` for the `ef_search` override limitation.
    fn filtered_search(
        &self,
        query: &[f32],
        k: usize,
        ef_search: usize,
        filter: &dyn Fn(VectorId) -> bool,
    ) -> Result<Vec<VectorSearchResult>, VectorError> {
        self.validate_dimensions(query)?;

        if ef_search != self.config.ef_search {
            tracing::warn!(
                requested = ef_search,
                default = self.config.ef_search,
                "per-query ef_search override not supported; using index-level default (may affect recall)"
            );
        }

        // Wrap the &dyn Fn in a concrete closure type that USearch's generic
        // filtered_search can monomorphize over.
        let predicate = |key: usearch::Key| filter(key);

        let matches = self
            .inner
            .filtered_search(query, k, predicate)
            .map_err(|e| VectorError::Backend(format!("USearch filtered_search failed: {e}")))?;

        Ok(matches
            .keys
            .into_iter()
            .zip(matches.distances)
            .map(|(id, distance)| VectorSearchResult { id, distance })
            .collect())
    }

    /// Delete a vector by ID. `USearch` uses lazy tombstoning internally --
    /// the vector slot is marked as removed but not reclaimed until a
    /// future compaction or rebuild.
    ///
    /// # Errors
    ///
    /// Returns [`VectorError::NotFound`] if no vector with this ID exists.
    fn delete(&self, id: VectorId) -> Result<(), VectorError> {
        if !self.inner.contains(id) {
            return Err(VectorError::NotFound { id });
        }

        let removed = self
            .inner
            .remove(id)
            .map_err(|e| VectorError::Backend(format!("USearch remove failed: {e}")))?;

        if removed == 0 {
            return Err(VectorError::NotFound { id });
        }

        // Do NOT decrement total_slots: the graph slot remains allocated
        // as a tombstone. total_slots tracks occupied graph slots (live +
        // tombstoned) so tombstone_ratio() can detect when compaction is needed.
        Ok(())
    }

    /// Pre-allocate capacity for `additional` more vectors.
    ///
    /// `USearch`'s `reserve()` takes a total capacity (not additional), so we
    /// add the current size to the requested additional count.
    fn reserve(&self, additional: usize) -> Result<(), VectorError> {
        let total = self.inner.size() + additional;
        self.inner
            .reserve(total)
            .map_err(|e| VectorError::Backend(format!("USearch reserve failed: {e}")))?;
        Ok(())
    }

    fn save(&self, path: &Path) -> Result<(), VectorError> {
        let path_str = Self::path_to_str(path)?;
        self.inner
            .save(path_str)
            .map_err(|e| VectorError::Backend(format!("USearch save failed: {e}")))?;
        Ok(())
    }

    fn load(path: &Path, config: &VectorIndexConfig) -> Result<Self, VectorError> {
        let index = Self::new(config.clone())?;
        let path_str = Self::path_to_str(path)?;
        index
            .inner
            .load(path_str)
            .map_err(|e| VectorError::Backend(format!("USearch load failed: {e}")))?;

        // Reconstruct total_slots from the loaded index size.
        // After load, all entries are live (no tombstones in a freshly loaded index).
        index
            .total_slots
            .store(index.inner.size(), Ordering::Relaxed);

        Ok(index)
    }

    fn view(path: &Path, config: &VectorIndexConfig) -> Result<Self, VectorError> {
        let index = Self::new(config.clone())?;
        let path_str = Self::path_to_str(path)?;
        index
            .inner
            .view(path_str)
            .map_err(|e| VectorError::Backend(format!("USearch view failed: {e}")))?;

        // After view, all entries are live (no tombstones).
        index
            .total_slots
            .store(index.inner.size(), Ordering::Relaxed);

        Ok(index)
    }

    /// Total number of occupied graph slots, including tombstoned entries.
    ///
    /// `USearch`'s `size()` excludes tombstoned entries, so we use our own
    /// `total_slots` counter which includes them.
    fn len(&self) -> usize {
        self.total_slots.load(Ordering::Relaxed)
    }

    /// Number of live (non-tombstoned) vectors in the index.
    ///
    /// Delegates to `USearch`'s `size()`, which already excludes removed entries.
    fn len_live(&self) -> usize {
        self.inner.size()
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;

    fn default_config(dimensions: usize) -> VectorIndexConfig {
        VectorIndexConfig {
            dimensions,
            metric: DistanceMetric::L2,
            quantization: QuantizationLevel::F16,
            connectivity: 16,
            ef_construction: 200,
            ef_search: 200,
        }
    }

    #[test]
    fn usearch_new_is_empty() {
        let index = UsearchIndex::new(default_config(128)).unwrap();
        assert!(index.is_empty());
        assert_eq!(index.len(), 0);
        assert_eq!(index.len_live(), 0);
    }

    #[test]
    fn usearch_insert_and_len() {
        let config = default_config(4);
        let index = UsearchIndex::new(config).unwrap();
        index.reserve(10).unwrap();
        index.insert(1, &[1.0, 2.0, 3.0, 4.0]).unwrap();
        index.insert(2, &[5.0, 6.0, 7.0, 8.0]).unwrap();
        assert_eq!(index.len(), 2);
        assert_eq!(index.len_live(), 2);
        assert!(!index.is_empty());
    }

    #[test]
    fn usearch_dimension_mismatch_insert() {
        let config = default_config(4);
        let index = UsearchIndex::new(config).unwrap();
        let result = index.insert(1, &[1.0, 2.0]); // 2D into 4D index
        assert!(result.is_err());
        match result.unwrap_err() {
            VectorError::DimensionMismatch { expected, got } => {
                assert_eq!(expected, 4);
                assert_eq!(got, 2);
            }
            other => panic!("expected DimensionMismatch, got {other:?}"),
        }
    }

    #[test]
    fn usearch_dimension_mismatch_search() {
        let config = default_config(4);
        let index = UsearchIndex::new(config).unwrap();
        index.reserve(2).unwrap();
        index.insert(1, &[1.0, 2.0, 3.0, 4.0]).unwrap();
        let result = index.search(&[1.0, 2.0], 1, 200); // 2D query on 4D index
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            VectorError::DimensionMismatch {
                expected: 4,
                got: 2
            }
        ));
    }

    #[test]
    fn usearch_delete_not_found() {
        let config = default_config(4);
        let index = UsearchIndex::new(config).unwrap();
        let result = index.delete(999);
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            VectorError::NotFound { id: 999 }
        ));
    }

    #[test]
    fn usearch_search_empty_index() {
        let config = default_config(4);
        let index = UsearchIndex::new(config).unwrap();
        let results = index.search(&[1.0, 2.0, 3.0, 4.0], 5, 200).unwrap();
        assert!(results.is_empty());
    }

    #[test]
    fn usearch_basic_search() {
        let config = default_config(4);
        let index = UsearchIndex::new(config).unwrap();
        index.reserve(10).unwrap();

        index.insert(1, &[1.0, 0.0, 0.0, 0.0]).unwrap();
        index.insert(2, &[0.0, 1.0, 0.0, 0.0]).unwrap();
        index.insert(3, &[0.0, 0.0, 1.0, 0.0]).unwrap();

        let results = index.search(&[1.0, 0.0, 0.0, 0.0], 3, 200).unwrap();
        assert_eq!(results.len(), 3);
        // Closest should be vector 1 (identical to query)
        assert_eq!(results[0].id, 1);
    }

    #[test]
    fn usearch_delete_and_search() {
        let config = default_config(4);
        let index = UsearchIndex::new(config).unwrap();
        index.reserve(10).unwrap();

        index.insert(1, &[1.0, 0.0, 0.0, 0.0]).unwrap();
        index.insert(2, &[0.0, 1.0, 0.0, 0.0]).unwrap();
        index.insert(3, &[0.0, 0.0, 1.0, 0.0]).unwrap();

        index.delete(2).unwrap();
        assert_eq!(index.len_live(), 2);

        let results = index.search(&[0.0, 1.0, 0.0, 0.0], 3, 200).unwrap();
        // Vector 2 must not appear in results
        assert!(results.iter().all(|r| r.id != 2));
    }

    #[test]
    fn usearch_is_send_and_sync() {
        fn assert_send_sync<T: Send + Sync>() {}
        assert_send_sync::<UsearchIndex>();
    }

    #[test]
    fn usearch_save_load_roundtrip() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("test_usearch.idx");

        let config = default_config(4);
        let index = UsearchIndex::new(config.clone()).unwrap();
        index.reserve(10).unwrap();
        index.insert(1, &[1.0, 0.0, 0.0, 0.0]).unwrap();
        index.insert(2, &[0.0, 1.0, 0.0, 0.0]).unwrap();

        index.save(&path).unwrap();

        let loaded = UsearchIndex::load(&path, &config).unwrap();
        assert_eq!(loaded.len(), 2);
        assert_eq!(loaded.len_live(), 2);

        let results = loaded.search(&[1.0, 0.0, 0.0, 0.0], 1, 200).unwrap();
        assert_eq!(results[0].id, 1);
    }

    #[test]
    fn usearch_view_readonly() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("test_usearch_view.idx");

        let config = default_config(4);
        let index = UsearchIndex::new(config.clone()).unwrap();
        index.reserve(10).unwrap();
        index.insert(1, &[1.0, 0.0, 0.0, 0.0]).unwrap();
        index.insert(2, &[0.0, 1.0, 0.0, 0.0]).unwrap();

        index.save(&path).unwrap();

        let viewed = UsearchIndex::view(&path, &config).unwrap();
        assert_eq!(viewed.len(), 2);

        let results = viewed.search(&[1.0, 0.0, 0.0, 0.0], 1, 200).unwrap();
        assert_eq!(results[0].id, 1);
    }

    #[test]
    fn usearch_filtered_search() {
        let config = default_config(4);
        let index = UsearchIndex::new(config).unwrap();
        index.reserve(10).unwrap();

        for i in 0..10_u64 {
            #[allow(clippy::cast_precision_loss)]
            let v = [i as f32 * 0.1, 0.0, 0.0, 0.0];
            index.insert(i, &v).unwrap();
        }

        // Only allow even IDs
        let results = index
            .filtered_search(&[0.5, 0.0, 0.0, 0.0], 5, 200, &|id| id % 2 == 0)
            .unwrap();

        for r in &results {
            assert!(
                r.id % 2 == 0,
                "odd ID {} in even-only filtered search",
                r.id
            );
        }
    }

    #[test]
    fn usearch_reserve() {
        let config = default_config(4);
        let index = UsearchIndex::new(config).unwrap();
        index.reserve(1000).unwrap();
        assert!(index.is_empty()); // no side effects
    }

    #[test]
    fn usearch_tombstone_ratio() {
        let config = default_config(4);
        let index = UsearchIndex::new(config).unwrap();
        index.reserve(10).unwrap();

        index.insert(1, &[1.0, 0.0, 0.0, 0.0]).unwrap();
        index.insert(2, &[0.0, 1.0, 0.0, 0.0]).unwrap();

        // No tombstones yet.
        assert_eq!(index.len(), 2);
        assert_eq!(index.len_live(), 2);
        assert!((index.tombstone_ratio() - 0.0).abs() < f64::EPSILON);

        index.delete(1).unwrap();
        // len() still shows 2 (total_slots includes the tombstoned slot).
        // len_live() shows 1 (USearch's size() excludes removed entries).
        // tombstone_ratio = (2 - 1) / 2 = 0.5
        assert_eq!(index.len(), 2);
        assert_eq!(index.len_live(), 1);
        assert!((index.tombstone_ratio() - 0.5).abs() < f64::EPSILON);
    }
}
