//! Vector storage backend trait abstraction
//!
//! Defines the `VectorBackend` trait that provides a unified interface
//! for vector storage and nearest-neighbor search across different backends:
//!
//! - **MmapVectorBackend**: Legacy memory-mapped binary files (always available)
//! - **LanceVectorBackend**: Lance columnar format (requires `lance-storage` feature)
//!
//! This trait replaces direct usage of `MmapStorage` and enables
//! transparent switching between vector storage strategies.

use std::io;

use super::backend::{StorageBackend, StorageError, StorageResult};

/// A single search result: (external_id, distance).
///
/// Lower distance = more similar (for L2/Cosine distance metrics).
#[derive(Debug, Clone)]
pub struct VectorSearchResult {
    /// External vector ID
    pub id: String,
    /// Distance to query vector
    pub distance: f32,
}

impl VectorSearchResult {
    pub fn new(id: impl Into<String>, distance: f32) -> Self {
        Self {
            id: id.into(),
            distance,
        }
    }
}

/// Configuration for creating a disk-based vector index (e.g., IVF_PQ, IVF_HNSW_SQ).
#[derive(Debug, Clone)]
pub struct VectorIndexConfig {
    /// Index type (e.g., IvfPq, IvfHnswSq)
    pub index_type: VectorIndexType,

    /// Number of partitions for IVF (default: sqrt(n) for IVF_PQ, sqrt(n)/10 for IVF_HNSW_*)
    pub num_partitions: Option<u32>,

    /// Number of sub-vectors for PQ (default: dim / 8). Only used for PQ-based indices.
    pub num_sub_vectors: Option<u32>,

    /// Distance metric
    pub metric: VectorDistanceMetric,

    // ---- HNSW parameters (for IvfHnswSq / IvfHnswPq / IvfHnswFlat) ----

    /// Max number of connections per node in the HNSW graph (default: 20).
    /// Higher values → better recall but larger index.
    pub hnsw_m: Option<usize>,

    /// ef_construction: number of candidates considered during graph build (default: 150).
    /// Higher → better graph quality but slower build time.
    pub hnsw_ef_construction: Option<usize>,

    /// Max level of the HNSW graph hierarchy (default: 7).
    pub hnsw_max_level: Option<u16>,

    // ---- SQ parameters (for IvfSq / IvfHnswSq) ----

    /// Number of bits for scalar quantization (default: 8). Valid: 4, 8.
    pub sq_num_bits: Option<u16>,

    /// Sample rate for SQ training (default: 256).
    pub sq_sample_rate: Option<usize>,

    // ---- PQ parameters ----

    /// Number of bits per centroid for PQ (default: 8). Currently only 8 is supported.
    pub pq_num_bits: Option<u8>,

    /// Max iterations for PQ kmeans training (default: 50).
    pub pq_max_iterations: Option<usize>,

    // ---- GPU acceleration parameters (optional, requires `gpu-index` feature) ----

    /// Enable GPU-accelerated index building.
    ///
    /// When `true` and the `gpu-index` feature is enabled, the system will:
    /// 1. Probe for a compatible GPU at runtime
    /// 2. If found, use GPU for KMeans training, quantization, and graph building
    /// 3. If not found, transparently fall back to CPU (zero impact)
    ///
    /// Default: `false` (CPU-only, for backward compatibility)
    pub gpu_accelerate: bool,

    /// GPU device ID for index building (-1 = auto-select best device).
    /// Only used when `gpu_accelerate` is `true`.
    pub gpu_device_id: Option<i32>,

    /// GPU memory budget in bytes for index building.
    /// `None` = use 80% of free GPU memory (recommended).
    pub gpu_memory_budget: Option<usize>,
}

impl Default for VectorIndexConfig {
    fn default() -> Self {
        Self {
            index_type: VectorIndexType::IvfPq,
            num_partitions: None,
            num_sub_vectors: None,
            metric: VectorDistanceMetric::Cosine,
            hnsw_m: None,
            hnsw_ef_construction: None,
            hnsw_max_level: None,
            sq_num_bits: None,
            sq_sample_rate: None,
            pq_num_bits: None,
            pq_max_iterations: None,
            gpu_accelerate: false,
            gpu_device_id: None,
            gpu_memory_budget: None,
        }
    }
}

impl VectorIndexConfig {
    /// Create a config for IVF_HNSW_SQ (DiskANN-style) index with sensible defaults.
    ///
    /// This is the recommended index type for large-scale vector search with
    /// disk-friendly access patterns. Uses HNSW graph for routing + SQ for
    /// memory-efficient storage.
    pub fn diskann(metric: VectorDistanceMetric) -> Self {
        Self {
            index_type: VectorIndexType::IvfHnswSq,
            metric,
            ..Default::default()
        }
    }

    /// Create a config for IVF_HNSW_SQ with explicit HNSW parameters.
    pub fn diskann_with_params(
        metric: VectorDistanceMetric,
        num_partitions: u32,
        hnsw_m: usize,
        ef_construction: usize,
    ) -> Self {
        Self {
            index_type: VectorIndexType::IvfHnswSq,
            metric,
            num_partitions: Some(num_partitions),
            hnsw_m: Some(hnsw_m),
            hnsw_ef_construction: Some(ef_construction),
            ..Default::default()
        }
    }

    /// Create a config for IVF_PQ index.
    pub fn ivf_pq(metric: VectorDistanceMetric) -> Self {
        Self {
            index_type: VectorIndexType::IvfPq,
            metric,
            ..Default::default()
        }
    }

    /// Create a config for IVF_SQ index (scalar quantization, no HNSW).
    pub fn ivf_sq(metric: VectorDistanceMetric) -> Self {
        Self {
            index_type: VectorIndexType::IvfSq,
            metric,
            ..Default::default()
        }
    }

    /// Create a config for IVF_RQ (binary/rabit quantization) index.
    ///
    /// Best for extreme compression scenarios where memory is critical.
    /// Uses 1-bit binary quantization for 32x compression.
    pub fn ivf_rq(metric: VectorDistanceMetric) -> Self {
        Self {
            index_type: VectorIndexType::IvfRq,
            metric,
            ..Default::default()
        }
    }

    /// Create a recommended config based on dataset characteristics.
    ///
    /// Automatically selects the best index type and quantization strategy
    /// based on the number of vectors, dimension, and optional memory budget.
    pub fn auto(
        metric: VectorDistanceMetric,
        num_vectors: usize,
        dimension: usize,
        target_memory_mb: Option<usize>,
    ) -> Self {
        let (index_type, _strategy) = QuantizationStrategy::recommend(
            num_vectors, dimension, target_memory_mb,
        );
        Self {
            index_type,
            metric,
            ..Default::default()
        }
    }

    /// Create a config for CAGRA graph index (GPU-exclusive).
    ///
    /// CAGRA is a GPU-native graph index that provides 10-100x faster build
    /// and search compared to HNSW on GPU hardware.
    ///
    /// If GPU is not available at runtime, automatically falls back to IVF_HNSW_SQ.
    pub fn cagra(metric: VectorDistanceMetric) -> Self {
        Self {
            index_type: VectorIndexType::Cagra,
            metric,
            gpu_accelerate: true,
            ..Default::default()
        }
    }

    /// Enable GPU acceleration on an existing config.
    ///
    /// When GPU is available, uses GPU for KMeans training and quantization.
    /// When GPU is not available, silently falls back to CPU (no error).
    pub fn with_gpu(mut self) -> Self {
        self.gpu_accelerate = true;
        self
    }

    /// Enable GPU acceleration with a specific device.
    pub fn with_gpu_device(mut self, device_id: i32) -> Self {
        self.gpu_accelerate = true;
        self.gpu_device_id = Some(device_id);
        self
    }

    /// Set GPU memory budget.
    pub fn with_gpu_memory_budget(mut self, budget_bytes: usize) -> Self {
        self.gpu_memory_budget = Some(budget_bytes);
        self
    }
}

/// Vector index type
///
/// Covers all index types supported by the Lance backend. Indices with
/// HNSW sub-graph (DiskANN-style) are recommended for datasets > 100k vectors,
/// as they provide better recall with fewer disk reads.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VectorIndexType {
    // ---- Flat (brute-force) ----

    /// Flat brute-force (exact NN, good for small datasets < 10k)
    Flat,

    // ---- IVF family ----

    /// IVF + Flat (no quantization, moderate-size datasets)
    IvfFlat,

    /// IVF + Scalar Quantization (memory-efficient, moderate recall)
    IvfSq,

    /// IVF + Product Quantization (good compression for large datasets)
    IvfPq,

    /// IVF + Rabit (Binary) Quantization (extreme compression, fast distance)
    ///
    /// Uses 1-bit binary quantization for maximum compression.
    /// Best for very large datasets where memory is critical.
    IvfRq,

    // ---- IVF + HNSW family (DiskANN-style, recommended for production) ----

    /// IVF + HNSW + Flat (best recall, highest memory usage)
    IvfHnswFlat,

    /// IVF + HNSW + Scalar Quantization (recommended: good recall + low memory)
    ///
    /// This is the "DiskANN" style index — HNSW graph for partition routing,
    /// SQ for compressed in-partition storage. Excellent disk-friendly access.
    IvfHnswSq,

    /// IVF + HNSW + Product Quantization (highest compression, good recall)
    IvfHnswPq,

    // ---- GPU-exclusive index types (require `gpu-index` feature) ----

    /// CAGRA graph index (GPU-exclusive, NVIDIA 2023).
    ///
    /// Compressed Adaptive Graph for Rapid Approximation — a GPU-native graph
    /// index that is 10-100x faster than HNSW for both build and search on GPU.
    /// Falls back to IVF_HNSW_SQ if GPU is not available at runtime.
    ///
    /// Requires: `gpu-index` feature + NVIDIA GPU with CC >= 6.0.
    Cagra,
}

impl std::fmt::Display for VectorIndexType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Flat => write!(f, "FLAT"),
            Self::IvfFlat => write!(f, "IVF_FLAT"),
            Self::IvfSq => write!(f, "IVF_SQ"),
            Self::IvfPq => write!(f, "IVF_PQ"),
            Self::IvfRq => write!(f, "IVF_RQ"),
            Self::IvfHnswFlat => write!(f, "IVF_HNSW_FLAT"),
            Self::IvfHnswSq => write!(f, "IVF_HNSW_SQ"),
            Self::IvfHnswPq => write!(f, "IVF_HNSW_PQ"),
            Self::Cagra => write!(f, "CAGRA"),
        }
    }
}

impl VectorIndexType {
    /// Get the quantization strategy used by this index type.
    pub fn quantization(&self) -> QuantizationStrategy {
        match self {
            Self::Flat | Self::IvfFlat | Self::IvfHnswFlat => QuantizationStrategy::None,
            Self::IvfSq | Self::IvfHnswSq => QuantizationStrategy::Scalar,
            Self::IvfPq | Self::IvfHnswPq => QuantizationStrategy::Product,
            Self::IvfRq => QuantizationStrategy::Binary,
            Self::Cagra => QuantizationStrategy::None,
        }
    }

    /// Whether this index type uses HNSW graph for routing (DiskANN-style).
    pub fn uses_hnsw(&self) -> bool {
        matches!(self, Self::IvfHnswFlat | Self::IvfHnswSq | Self::IvfHnswPq)
    }

    /// Whether this index type uses IVF partitioning.
    pub fn uses_ivf(&self) -> bool {
        !matches!(self, Self::Flat | Self::Cagra)
    }

    /// Whether this index type is GPU-exclusive.
    ///
    /// GPU-exclusive index types require the `gpu-index` feature and a compatible
    /// NVIDIA GPU. When GPU is not available, they fall back to a CPU-compatible
    /// alternative (e.g., CAGRA → IVF_HNSW_SQ).
    pub fn is_gpu_exclusive(&self) -> bool {
        matches!(self, Self::Cagra)
    }

    /// Get the CPU fallback index type for GPU-exclusive indices.
    ///
    /// Returns `None` for CPU-compatible indices (they don't need a fallback).
    pub fn cpu_fallback(&self) -> Option<Self> {
        match self {
            Self::Cagra => Some(Self::IvfHnswSq),
            _ => None,
        }
    }
}

// ==========================================================================
// Quantization Strategy
// ==========================================================================

/// Quantization strategy for vector compression.
///
/// Quantization reduces memory usage and speeds up distance computation
/// by compressing vectors into lower-precision representations.
///
/// # Comparison
///
/// | Strategy | Bits/dim | Compression | Recall | Speed |
/// |----------|----------|-------------|--------|-------|
/// | None     | 32       | 1x          | Best   | Base  |
/// | Scalar   | 8        | 4x          | Good   | Fast  |
/// | Product  | ~2-4     | 8-16x       | Good   | Fast  |
/// | Binary   | 1        | 32x         | Lower  | Fastest |
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum QuantizationStrategy {
    /// No quantization — full float32 precision.
    None,

    /// Scalar Quantization (SQ): each dimension quantized to uint8.
    ///
    /// Provides 4x compression with minimal recall loss.
    /// Best balance of compression and accuracy.
    Scalar,

    /// Product Quantization (PQ): vectors split into sub-vectors,
    /// each encoded as a centroid index.
    ///
    /// Provides 8-16x compression. Good for very large datasets.
    Product,

    /// Binary Quantization (BQ/RQ): each dimension quantized to 1 bit.
    ///
    /// Provides 32x compression and fastest hamming distance.
    /// Best for coarse filtering, often combined with re-ranking.
    Binary,
}

impl std::fmt::Display for QuantizationStrategy {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::None => write!(f, "None"),
            Self::Scalar => write!(f, "SQ (Scalar)"),
            Self::Product => write!(f, "PQ (Product)"),
            Self::Binary => write!(f, "BQ (Binary)"),
        }
    }
}

impl QuantizationStrategy {
    /// Estimated compression ratio compared to float32.
    ///
    /// Returns the approximate factor by which memory usage is reduced.
    /// For example, SQ returns 4.0 (4x smaller than float32).
    pub fn compression_ratio(&self) -> f64 {
        match self {
            Self::None => 1.0,
            Self::Scalar => 4.0,     // 32-bit → 8-bit
            Self::Product => 12.0,   // typical PQ with 8-bit codes, ~2-3 bits/dim
            Self::Binary => 32.0,    // 32-bit → 1-bit
        }
    }

    /// Estimate memory usage in bytes for a vector dataset.
    ///
    /// # Arguments
    /// * `num_vectors` - Number of vectors in the dataset
    /// * `dimension` - Dimension of each vector
    ///
    /// # Returns
    /// Estimated memory in bytes (approximate, excludes index overhead)
    pub fn estimate_memory_bytes(&self, num_vectors: usize, dimension: usize) -> usize {
        let raw_bytes = num_vectors * dimension * 4; // float32 = 4 bytes
        (raw_bytes as f64 / self.compression_ratio()) as usize
    }

    /// Estimate memory usage as a human-readable string.
    pub fn estimate_memory_human(&self, num_vectors: usize, dimension: usize) -> String {
        let bytes = self.estimate_memory_bytes(num_vectors, dimension);
        if bytes >= 1 << 30 {
            format!("{:.2} GB", bytes as f64 / (1u64 << 30) as f64)
        } else if bytes >= 1 << 20 {
            format!("{:.2} MB", bytes as f64 / (1u64 << 20) as f64)
        } else if bytes >= 1 << 10 {
            format!("{:.2} KB", bytes as f64 / (1u64 << 10) as f64)
        } else {
            format!("{} B", bytes)
        }
    }

    /// Recommend a quantization strategy based on dataset characteristics.
    ///
    /// # Arguments
    /// * `num_vectors` - Expected number of vectors
    /// * `dimension` - Vector dimension
    /// * `target_memory_mb` - Optional memory budget in MB
    ///
    /// # Returns
    /// Recommended `(VectorIndexType, QuantizationStrategy)` pair
    pub fn recommend(
        num_vectors: usize,
        dimension: usize,
        target_memory_mb: Option<usize>,
    ) -> (VectorIndexType, QuantizationStrategy) {
        let raw_mb = (num_vectors * dimension * 4) / (1 << 20);

        // If there's a memory budget, pick the smallest strategy that fits
        if let Some(budget) = target_memory_mb {
            if raw_mb <= budget {
                // Fits without compression → use HNSW + Flat for best recall
                if num_vectors > 100_000 {
                    return (VectorIndexType::IvfHnswFlat, QuantizationStrategy::None);
                } else {
                    return (VectorIndexType::IvfFlat, QuantizationStrategy::None);
                }
            }
            let sq_mb = raw_mb / 4;
            if sq_mb <= budget {
                return if num_vectors > 100_000 {
                    (VectorIndexType::IvfHnswSq, QuantizationStrategy::Scalar)
                } else {
                    (VectorIndexType::IvfSq, QuantizationStrategy::Scalar)
                };
            }
            let pq_mb = raw_mb / 12;
            if pq_mb <= budget {
                return if num_vectors > 100_000 {
                    (VectorIndexType::IvfHnswPq, QuantizationStrategy::Product)
                } else {
                    (VectorIndexType::IvfPq, QuantizationStrategy::Product)
                };
            }
            // Even PQ doesn't fit → use Binary
            return (VectorIndexType::IvfRq, QuantizationStrategy::Binary);
        }

        // No memory budget — pick based on dataset size
        match num_vectors {
            0..=10_000 => (VectorIndexType::Flat, QuantizationStrategy::None),
            10_001..=100_000 => (VectorIndexType::IvfPq, QuantizationStrategy::Product),
            100_001..=10_000_000 => (VectorIndexType::IvfHnswSq, QuantizationStrategy::Scalar),
            _ => (VectorIndexType::IvfHnswPq, QuantizationStrategy::Product),
        }
    }
}

/// Distance metric for vector similarity
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VectorDistanceMetric {
    /// Cosine distance (1 - cosine_similarity)
    Cosine,
    /// Euclidean (L2) distance
    L2,
    /// Dot product (negative, for max inner product search)
    Dot,
}

impl std::fmt::Display for VectorDistanceMetric {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Cosine => write!(f, "cosine"),
            Self::L2 => write!(f, "L2"),
            Self::Dot => write!(f, "dot"),
        }
    }
}

/// Unified vector storage and search trait.
///
/// Provides vector CRUD operations and nearest-neighbor search.
/// Implementations can use different backing stores (mmap, Lance, etc.)
/// and different index strategies (HNSW in-memory, IVF_PQ disk, brute-force).
///
/// # Thread Safety
///
/// All implementations must be `Send + Sync` for concurrent access.
pub trait VectorBackend: Send + Sync {
    // ====================================================================
    // Vector CRUD
    // ====================================================================

    /// Store a vector with the given external ID.
    ///
    /// If a vector with this ID already exists, it is overwritten.
    fn store_vector(&self, id: &str, vector: &[f32]) -> StorageResult<()>;

    /// Retrieve a vector by external ID.
    fn get_vector(&self, id: &str) -> StorageResult<Option<Vec<f32>>>;

    /// Retrieve multiple vectors by IDs.
    ///
    /// Default implementation calls `get_vector` in a loop.
    fn get_vectors_batch(&self, ids: &[&str]) -> StorageResult<Vec<Option<Vec<f32>>>> {
        ids.iter().map(|id| self.get_vector(id)).collect()
    }

    /// Delete a vector by ID.
    fn delete_vector(&self, id: &str) -> StorageResult<bool>;

    /// Check if a vector exists.
    fn contains_vector(&self, id: &str) -> StorageResult<bool>;

    // ====================================================================
    // Counting & Listing
    // ====================================================================

    /// Get the number of stored vectors.
    fn vector_count(&self) -> StorageResult<usize>;

    /// Get the vector dimension.
    fn dimension(&self) -> usize;

    /// List all vector IDs.
    fn vector_ids(&self) -> StorageResult<Vec<String>>;

    // ====================================================================
    // Search
    // ====================================================================

    /// Search for the k nearest neighbors to a query vector.
    ///
    /// Returns results sorted by distance (ascending).
    fn search_nearest(
        &self,
        query: &[f32],
        k: usize,
    ) -> StorageResult<Vec<VectorSearchResult>>;

    /// Search for nearest neighbors with a pre-filter on IDs.
    ///
    /// Only considers vectors whose IDs are in `allowed_ids`.
    /// Default implementation falls back to `search_nearest` then filters.
    fn search_nearest_filtered(
        &self,
        query: &[f32],
        k: usize,
        allowed_ids: &[String],
    ) -> StorageResult<Vec<VectorSearchResult>> {
        // Default: search more candidates, then filter
        let candidates = self.search_nearest(query, k * 4)?;
        let mut filtered: Vec<VectorSearchResult> = candidates
            .into_iter()
            .filter(|r| allowed_ids.contains(&r.id))
            .collect();
        filtered.truncate(k);
        Ok(filtered)
    }

    /// Search for nearest neighbors with a `FilterExpr` pre-filter.
    ///
    /// For Lance backend, this translates to a Lance scanner `filter()` call
    /// that combines vector search with scalar predicate push-down — achieving
    /// true **Filtered ANN** with efficient index usage.
    ///
    /// Default implementation: evaluates `FilterExpr` on doc store, gets allowed
    /// IDs, then calls `search_nearest_filtered`.
    fn search_nearest_with_expr(
        &self,
        query: &[f32],
        k: usize,
        expr: &crate::storage::filter_expr::FilterExpr,
        doc_store: Option<&dyn StorageBackend>,
    ) -> StorageResult<Vec<VectorSearchResult>> {
        // Default: if we have a doc store, filter IDs first, then search
        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_nearest_filtered(query, k, &allowed_ids)
        } else {
            // No doc store — fall back to unfiltered search
            self.search_nearest(query, k)
        }
    }

    // ====================================================================
    // Index Management
    // ====================================================================

    /// Create a vector index for accelerated search.
    ///
    /// Not all backends support this — backends that don't (e.g., mmap + HNSW)
    /// may return `Ok(())` as a no-op.
    fn create_index(&self, _config: &VectorIndexConfig) -> StorageResult<()> {
        Ok(()) // default: no-op
    }

    /// Check if a vector index exists and is up-to-date.
    fn has_index(&self) -> StorageResult<bool> {
        Ok(false)
    }

    /// Get information about the current index (if any).
    ///
    /// Returns `None` if no index exists.
    fn index_info(&self) -> StorageResult<Option<VectorIndexInfo>> {
        Ok(None) // default: no index info available
    }

    // ====================================================================
    // Persistence
    // ====================================================================

    /// Flush any pending writes to durable storage.
    fn flush(&self) -> StorageResult<()>;

    // ====================================================================
    // Backend Info
    // ====================================================================

    /// Backend type name (e.g., "mmap", "lance-vector").
    fn backend_name(&self) -> &str;
}

// ==========================================================================
// Vector Index Info
// ==========================================================================

/// Information about an existing vector index.
///
/// Provides metadata about the current index configuration,
/// quantization strategy, and estimated memory usage.
#[derive(Debug, Clone)]
pub struct VectorIndexInfo {
    /// The index type.
    pub index_type: VectorIndexType,
    /// The quantization strategy.
    pub quantization: QuantizationStrategy,
    /// Number of indexed vectors.
    pub num_vectors: usize,
    /// Vector dimension.
    pub dimension: usize,
    /// Number of index segments/partitions.
    pub num_partitions: Option<usize>,
    /// Estimated index size in bytes.
    pub estimated_size_bytes: usize,
    /// Distance metric used.
    pub metric: VectorDistanceMetric,
}

impl VectorIndexInfo {
    /// Get estimated memory as a human-readable string.
    pub fn estimated_size_human(&self) -> String {
        let bytes = self.estimated_size_bytes;
        if bytes >= 1 << 30 {
            format!("{:.2} GB", bytes as f64 / (1u64 << 30) as f64)
        } else if bytes >= 1 << 20 {
            format!("{:.2} MB", bytes as f64 / (1u64 << 20) as f64)
        } else if bytes >= 1 << 10 {
            format!("{:.2} KB", bytes as f64 / (1u64 << 10) as f64)
        } else {
            format!("{} B", bytes)
        }
    }

    /// Get estimated compression ratio vs raw float32.
    pub fn compression_ratio(&self) -> f64 {
        let raw_bytes = self.num_vectors * self.dimension * 4;
        if self.estimated_size_bytes > 0 {
            raw_bytes as f64 / self.estimated_size_bytes as f64
        } else {
            1.0
        }
    }
}

// ==========================================================================
// Multi-Vector Support
// ==========================================================================

/// Configuration for a named vector column in a multi-vector dataset.
///
/// Each vector column represents a different embedding space (e.g.,
/// "title_embedding" with 384 dims and "content_embedding" with 1536 dims).
#[derive(Debug, Clone)]
pub struct VectorColumnConfig {
    /// Column name in the dataset (e.g., "title_embedding").
    pub name: String,
    /// Dimension of vectors in this column.
    pub dimension: usize,
    /// Distance metric for search on this column.
    pub metric: VectorDistanceMetric,
}

impl VectorColumnConfig {
    pub fn new(name: impl Into<String>, dimension: usize) -> Self {
        Self {
            name: name.into(),
            dimension,
            metric: VectorDistanceMetric::Cosine,
        }
    }

    pub fn with_metric(mut self, metric: VectorDistanceMetric) -> Self {
        self.metric = metric;
        self
    }
}

/// Strategy for fusing results from multiple vector columns.
///
/// When searching across multiple embedding columns, results need to
/// be combined into a single ranked list.
#[derive(Debug, Clone)]
pub enum MultiVectorFusion {
    /// Take the minimum distance from any column (best single match).
    ///
    /// Good when any single column match is sufficient.
    MinDistance,

    /// Weighted average of distances across columns.
    ///
    /// Weights correspond to column order. Missing vectors in a column
    /// are skipped (their weight is redistributed).
    WeightedAverage(Vec<f32>),

    /// Reciprocal Rank Fusion (RRF) — merges ranked lists by position.
    ///
    /// `k` is the smoothing parameter (default: 60.0).
    /// Good when columns have different distance scales.
    RRF { k: f32 },
}

impl Default for MultiVectorFusion {
    fn default() -> Self {
        Self::RRF { k: 60.0 }
    }
}

/// Multi-vector backend supporting multiple embedding columns per record.
///
/// Extends `VectorBackend` with column-aware operations. The base
/// `VectorBackend` methods operate on the "primary" column (the first
/// or default column).
///
/// # Use Cases
///
/// - Title + content embeddings for documents
/// - Image + text embeddings for multimodal data
/// - Multiple model embeddings for ensemble search
///
/// # Example
///
/// ```ignore
/// // Configure two vector columns
/// let columns = vec![
///     VectorColumnConfig::new("title_embedding", 384),
///     VectorColumnConfig::new("content_embedding", 1536),
/// ];
///
/// // Store vectors for each column
/// backend.store_vector_column("doc1", "title_embedding", &title_vec)?;
/// backend.store_vector_column("doc1", "content_embedding", &content_vec)?;
///
/// // Search across both columns with RRF fusion
/// let results = backend.search_multi_column(
///     &[("title_embedding", &query_384d), ("content_embedding", &query_1536d)],
///     10,
///     MultiVectorFusion::RRF { k: 60.0 },
/// )?;
/// ```
pub trait MultiVectorBackend: VectorBackend {
    /// List all configured vector columns.
    fn vector_columns(&self) -> Vec<VectorColumnConfig>;

    /// Store a vector in a specific column.
    fn store_vector_column(
        &self,
        id: &str,
        column: &str,
        vector: &[f32],
    ) -> StorageResult<()>;

    /// Get a vector from a specific column.
    fn get_vector_column(
        &self,
        id: &str,
        column: &str,
    ) -> StorageResult<Option<Vec<f32>>>;

    /// Delete a vector from a specific column.
    fn delete_vector_column(
        &self,
        id: &str,
        column: &str,
    ) -> StorageResult<bool>;

    /// Search nearest neighbors in a specific column.
    fn search_nearest_column(
        &self,
        column: &str,
        query: &[f32],
        k: usize,
    ) -> StorageResult<Vec<VectorSearchResult>>;

    /// Create an index for a specific column.
    fn create_index_column(
        &self,
        column: &str,
        config: &VectorIndexConfig,
    ) -> StorageResult<()>;

    /// Fused multi-column search — search multiple columns and merge results.
    ///
    /// Each entry in `queries` is a `(column_name, query_vector)` pair.
    /// Results from all columns are combined using the specified fusion strategy.
    fn search_multi_column(
        &self,
        queries: &[(&str, &[f32])],
        k: usize,
        fusion: MultiVectorFusion,
    ) -> StorageResult<Vec<VectorSearchResult>> {
        // Default implementation: search each column, then fuse
        let mut per_column_results: Vec<Vec<VectorSearchResult>> = Vec::new();

        for (col, query) in queries {
            let results = self.search_nearest_column(col, query, k * 2)?;
            per_column_results.push(results);
        }

        fuse_multi_column_results(&per_column_results, k, &fusion)
    }
}

/// Fuse results from multiple column searches into a single ranked list.
pub(crate) fn fuse_multi_column_results(
    column_results: &[Vec<VectorSearchResult>],
    k: usize,
    fusion: &MultiVectorFusion,
) -> StorageResult<Vec<VectorSearchResult>> {
    use std::collections::HashMap as FuseMap;

    if column_results.is_empty() {
        return Ok(Vec::new());
    }

    match fusion {
        MultiVectorFusion::MinDistance => {
            // For each ID, take the minimum distance across all columns
            let mut best: FuseMap<String, f32> = FuseMap::new();
            for results in column_results {
                for r in results {
                    let entry = best.entry(r.id.clone()).or_insert(f32::MAX);
                    if r.distance < *entry {
                        *entry = r.distance;
                    }
                }
            }
            let mut fused: Vec<VectorSearchResult> = best
                .into_iter()
                .map(|(id, dist)| VectorSearchResult::new(id, dist))
                .collect();
            fused.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap_or(std::cmp::Ordering::Equal));
            fused.truncate(k);
            Ok(fused)
        }

        MultiVectorFusion::WeightedAverage(weights) => {
            // Weighted average of distances
            let mut sums: FuseMap<String, (f64, f64)> = FuseMap::new(); // (weighted_dist_sum, weight_sum)
            for (i, results) in column_results.iter().enumerate() {
                let w = weights.get(i).copied().unwrap_or(1.0) as f64;
                for r in results {
                    let entry = sums.entry(r.id.clone()).or_insert((0.0, 0.0));
                    entry.0 += r.distance as f64 * w;
                    entry.1 += w;
                }
            }
            let mut fused: Vec<VectorSearchResult> = sums
                .into_iter()
                .map(|(id, (weighted_sum, weight_total))| {
                    let avg_dist = if weight_total > 0.0 {
                        (weighted_sum / weight_total) as f32
                    } else {
                        f32::MAX
                    };
                    VectorSearchResult::new(id, avg_dist)
                })
                .collect();
            fused.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap_or(std::cmp::Ordering::Equal));
            fused.truncate(k);
            Ok(fused)
        }

        MultiVectorFusion::RRF { k: rrf_k } => {
            // Reciprocal Rank Fusion: score = Σ 1/(rank + k) across columns
            let mut scores: FuseMap<String, f64> = FuseMap::new();
            for results in column_results {
                for (rank, r) in results.iter().enumerate() {
                    let rrf_score = 1.0 / (rank as f64 + *rrf_k as f64);
                    *scores.entry(r.id.clone()).or_insert(0.0) += rrf_score;
                }
            }
            // Higher RRF score = better → convert to distance (lower = better)
            // by using 1/score as distance
            let mut fused: Vec<VectorSearchResult> = scores
                .into_iter()
                .map(|(id, score)| {
                    let dist = if score > 0.0 { 1.0 / score as f32 } else { f32::MAX };
                    VectorSearchResult::new(id, dist)
                })
                .collect();
            fused.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap_or(std::cmp::Ordering::Equal));
            fused.truncate(k);
            Ok(fused)
        }
    }
}

// =========================================================================
// MmapVectorBackend — wraps legacy MmapStorage
// =========================================================================

use super::mmap::MmapStorage;
use crate::types::StorageConfig;
use crate::vector::distance::{Distance, DistanceCalculator};
use crate::types::DistanceMetric;

/// Legacy vector backend using memory-mapped binary files.
///
/// Wraps `MmapStorage` and adds brute-force nearest-neighbor search.
/// Suitable for small-to-medium datasets where HNSW is used as the
/// primary search index.
pub struct MmapVectorBackend {
    storage: MmapStorage,
    distance_calc: DistanceCalculator,
}

impl MmapVectorBackend {
    /// Create or open an MmapVectorBackend.
    pub fn open_or_create(config: StorageConfig, dimension: usize) -> io::Result<Self> {
        let storage = MmapStorage::open_or_create(config, dimension)?;
        Ok(Self {
            storage,
            distance_calc: DistanceCalculator::new(DistanceMetric::Cosine),
        })
    }

    /// Get a reference to the underlying `MmapStorage`.
    pub fn inner(&self) -> &MmapStorage {
        &self.storage
    }
}

impl VectorBackend for MmapVectorBackend {
    fn store_vector(&self, id: &str, vector: &[f32]) -> StorageResult<()> {
        self.storage.store(id, vector)?;
        Ok(())
    }

    fn get_vector(&self, id: &str) -> StorageResult<Option<Vec<f32>>> {
        match self.storage.get(id) {
            Ok(v) => Ok(Some(v)),
            Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(None),
            Err(e) => Err(StorageError::Io(e)),
        }
    }

    fn delete_vector(&self, _id: &str) -> StorageResult<bool> {
        // MmapStorage doesn't support deletion natively
        // (would require compaction). Return false.
        Ok(false)
    }

    fn contains_vector(&self, id: &str) -> StorageResult<bool> {
        Ok(self.storage.contains(id))
    }

    fn vector_count(&self) -> StorageResult<usize> {
        Ok(self.storage.len())
    }

    fn dimension(&self) -> usize {
        self.storage.dimension()
    }

    fn vector_ids(&self) -> StorageResult<Vec<String>> {
        Ok(self.storage.ids())
    }

    fn search_nearest(
        &self,
        query: &[f32],
        k: usize,
    ) -> StorageResult<Vec<VectorSearchResult>> {
        // Brute-force search over all stored vectors
        let ids = self.storage.ids();
        let mut results: Vec<VectorSearchResult> = Vec::with_capacity(ids.len());

        for id in &ids {
            if let Ok(vector) = self.storage.get(id) {
                let dist = self.distance_calc.calculate(query, &vector);
                results.push(VectorSearchResult::new(id.clone(), dist as f32));
            }
        }

        // Sort by distance (ascending)
        results.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap_or(std::cmp::Ordering::Equal));
        results.truncate(k);

        Ok(results)
    }

    fn flush(&self) -> StorageResult<()> {
        self.storage.flush()?;
        Ok(())
    }

    fn backend_name(&self) -> &str {
        "mmap"
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    #[test]
    fn test_mmap_vector_backend_basic() {
        let dir = tempdir().unwrap();
        let config = StorageConfig {
            path: dir.path().to_string_lossy().to_string(),
            ..Default::default()
        };

        let backend = MmapVectorBackend::open_or_create(config, 4).unwrap();
        assert_eq!(backend.dimension(), 4);
        assert_eq!(backend.vector_count().unwrap(), 0);
        assert_eq!(backend.backend_name(), "mmap");

        // Store vectors
        backend.store_vector("v1", &[1.0, 0.0, 0.0, 0.0]).unwrap();
        backend.store_vector("v2", &[0.0, 1.0, 0.0, 0.0]).unwrap();
        backend.store_vector("v3", &[0.9, 0.1, 0.0, 0.0]).unwrap();

        assert_eq!(backend.vector_count().unwrap(), 3);
        assert!(backend.contains_vector("v1").unwrap());
        assert!(!backend.contains_vector("v999").unwrap());

        // Get vector
        let v = backend.get_vector("v1").unwrap().unwrap();
        assert_eq!(v, vec![1.0, 0.0, 0.0, 0.0]);

        // Search nearest to [1,0,0,0] — should return v1 first, then v3
        let results = backend.search_nearest(&[1.0, 0.0, 0.0, 0.0], 2).unwrap();
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].id, "v1");
        assert!(results[0].distance < results[1].distance);
    }

    #[test]
    fn test_mmap_vector_backend_batch() {
        let dir = tempdir().unwrap();
        let config = StorageConfig {
            path: dir.path().to_string_lossy().to_string(),
            ..Default::default()
        };

        let backend = MmapVectorBackend::open_or_create(config, 3).unwrap();

        for i in 0..5 {
            backend.store_vector(&format!("vec_{}", i), &[i as f32, 0.0, 0.0]).unwrap();
        }

        let ids: Vec<&str> = vec!["vec_0", "vec_2", "vec_4", "vec_99"];
        let batch = backend.get_vectors_batch(&ids).unwrap();
        assert!(batch[0].is_some());
        assert!(batch[1].is_some());
        assert!(batch[2].is_some());
        assert!(batch[3].is_none()); // vec_99 doesn't exist
    }

    // ---- QuantizationStrategy tests ----

    #[test]
    fn test_quantization_compression_ratio() {
        assert_eq!(QuantizationStrategy::None.compression_ratio(), 1.0);
        assert_eq!(QuantizationStrategy::Scalar.compression_ratio(), 4.0);
        assert_eq!(QuantizationStrategy::Product.compression_ratio(), 12.0);
        assert_eq!(QuantizationStrategy::Binary.compression_ratio(), 32.0);
    }

    #[test]
    fn test_quantization_estimate_memory() {
        let num_vectors = 1_000_000;
        let dim = 1536;

        // Raw: 1M * 1536 * 4 = 5.72 GB
        let raw = QuantizationStrategy::None.estimate_memory_bytes(num_vectors, dim);
        assert_eq!(raw, num_vectors * dim * 4);

        // SQ: ~4x smaller
        let sq = QuantizationStrategy::Scalar.estimate_memory_bytes(num_vectors, dim);
        assert!(sq < raw);
        assert_eq!(sq, raw / 4);

        // PQ: ~12x smaller
        let pq = QuantizationStrategy::Product.estimate_memory_bytes(num_vectors, dim);
        assert!(pq < sq);

        // BQ: ~32x smaller
        let bq = QuantizationStrategy::Binary.estimate_memory_bytes(num_vectors, dim);
        assert!(bq < pq);

        // Human-readable
        let human = QuantizationStrategy::None.estimate_memory_human(num_vectors, dim);
        assert!(human.contains("GB"));
    }

    #[test]
    fn test_quantization_recommend_small() {
        let (idx_type, quant) = QuantizationStrategy::recommend(5_000, 128, None);
        assert_eq!(idx_type, VectorIndexType::Flat);
        assert_eq!(quant, QuantizationStrategy::None);
    }

    #[test]
    fn test_quantization_recommend_medium() {
        let (idx_type, quant) = QuantizationStrategy::recommend(50_000, 768, None);
        assert_eq!(idx_type, VectorIndexType::IvfPq);
        assert_eq!(quant, QuantizationStrategy::Product);
    }

    #[test]
    fn test_quantization_recommend_large() {
        let (idx_type, quant) = QuantizationStrategy::recommend(1_000_000, 1536, None);
        assert_eq!(idx_type, VectorIndexType::IvfHnswSq);
        assert_eq!(quant, QuantizationStrategy::Scalar);
    }

    #[test]
    fn test_quantization_recommend_with_budget() {
        // 1M vectors * 1536 dims * 4 bytes = ~5.7 GB raw
        // Budget: 500 MB → should pick PQ or tighter
        let (idx_type, quant) = QuantizationStrategy::recommend(1_000_000, 1536, Some(500));
        assert_eq!(quant, QuantizationStrategy::Product);
        assert!(idx_type.uses_hnsw()); // should use HNSW for 1M vectors
    }

    #[test]
    fn test_quantization_recommend_tight_budget() {
        // Very tight budget → should pick Binary
        let (idx_type, quant) = QuantizationStrategy::recommend(10_000_000, 1536, Some(100));
        assert_eq!(quant, QuantizationStrategy::Binary);
        assert_eq!(idx_type, VectorIndexType::IvfRq);
    }

    // ---- VectorIndexType tests ----

    #[test]
    fn test_index_type_quantization() {
        assert_eq!(VectorIndexType::Flat.quantization(), QuantizationStrategy::None);
        assert_eq!(VectorIndexType::IvfFlat.quantization(), QuantizationStrategy::None);
        assert_eq!(VectorIndexType::IvfSq.quantization(), QuantizationStrategy::Scalar);
        assert_eq!(VectorIndexType::IvfPq.quantization(), QuantizationStrategy::Product);
        assert_eq!(VectorIndexType::IvfRq.quantization(), QuantizationStrategy::Binary);
        assert_eq!(VectorIndexType::IvfHnswFlat.quantization(), QuantizationStrategy::None);
        assert_eq!(VectorIndexType::IvfHnswSq.quantization(), QuantizationStrategy::Scalar);
        assert_eq!(VectorIndexType::IvfHnswPq.quantization(), QuantizationStrategy::Product);
    }

    #[test]
    fn test_index_type_uses_hnsw() {
        assert!(!VectorIndexType::Flat.uses_hnsw());
        assert!(!VectorIndexType::IvfPq.uses_hnsw());
        assert!(VectorIndexType::IvfHnswSq.uses_hnsw());
        assert!(VectorIndexType::IvfHnswPq.uses_hnsw());
        assert!(VectorIndexType::IvfHnswFlat.uses_hnsw());
    }

    #[test]
    fn test_index_type_display() {
        assert_eq!(format!("{}", VectorIndexType::IvfHnswSq), "IVF_HNSW_SQ");
        assert_eq!(format!("{}", VectorIndexType::IvfPq), "IVF_PQ");
        assert_eq!(format!("{}", VectorIndexType::IvfRq), "IVF_RQ");
        assert_eq!(format!("{}", VectorIndexType::Flat), "FLAT");
    }

    #[test]
    fn test_vector_index_config_constructors() {
        let diskann = VectorIndexConfig::diskann(VectorDistanceMetric::Cosine);
        assert_eq!(diskann.index_type, VectorIndexType::IvfHnswSq);
        assert_eq!(diskann.metric, VectorDistanceMetric::Cosine);

        let ivf_pq = VectorIndexConfig::ivf_pq(VectorDistanceMetric::L2);
        assert_eq!(ivf_pq.index_type, VectorIndexType::IvfPq);

        let custom = VectorIndexConfig::diskann_with_params(
            VectorDistanceMetric::Dot, 16, 32, 200,
        );
        assert_eq!(custom.hnsw_m, Some(32));
        assert_eq!(custom.hnsw_ef_construction, Some(200));
        assert_eq!(custom.num_partitions, Some(16));
    }

    // ---- VectorIndexInfo tests ----

    #[test]
    fn test_vector_index_info() {
        let info = VectorIndexInfo {
            index_type: VectorIndexType::IvfHnswSq,
            quantization: QuantizationStrategy::Scalar,
            num_vectors: 1_000_000,
            dimension: 1536,
            num_partitions: Some(32),
            estimated_size_bytes: 1_000_000 * 1536, // 1 byte per dim after SQ
            metric: VectorDistanceMetric::Cosine,
        };

        let human = info.estimated_size_human();
        assert!(human.contains("GB") || human.contains("MB"));

        let ratio = info.compression_ratio();
        assert!(ratio > 3.0 && ratio < 5.0); // ~4x for SQ
    }
}
