rlm_rs/search/

mod.rs

//! Hybrid search with semantic and lexical retrieval.
//!
//! Combines vector similarity search with FTS5 BM25 using Reciprocal Rank Fusion (RRF).
//!
//! ## Features
//!
//! - **Semantic Search**: Vector similarity using embeddings
//! - **BM25 Search**: Full-text search using `SQLite` `FTS5`
//! - **Hybrid Search**: Combines both using Reciprocal Rank Fusion
//! - **HNSW Index**: Optional scalable approximate nearest neighbor search (requires `usearch-hnsw` feature)

pub mod hnsw;
mod rrf;

pub use hnsw::{HnswConfig, HnswIndex, HnswResult};
pub use rrf::{RrfConfig, reciprocal_rank_fusion, weighted_rrf};

use crate::embedding::{Embedder, cosine_similarity};
use crate::error::Result;
use crate::storage::{SqliteStorage, Storage};

/// Default similarity threshold for semantic search.
pub const DEFAULT_SIMILARITY_THRESHOLD: f32 = 0.3;

/// Default number of results to return.
pub const DEFAULT_TOP_K: usize = 10;

/// Search result with chunk ID and combined score.
#[derive(Debug, Clone)]
pub struct SearchResult {
    /// Chunk ID.
    pub chunk_id: i64,
    /// Buffer ID this chunk belongs to.
    pub buffer_id: i64,
    /// Sequential index within the buffer (0-based, for temporal ordering).
    pub index: usize,
    /// Combined RRF score (higher is better).
    pub score: f64,
    /// Semantic similarity score (if available).
    pub semantic_score: Option<f32>,
    /// BM25 score (if available).
    pub bm25_score: Option<f64>,
    /// Content preview (first N characters, if requested).
    pub content_preview: Option<String>,
}

/// Configuration for hybrid search.
#[derive(Debug, Clone)]
pub struct SearchConfig {
    /// Maximum number of results to return.
    pub top_k: usize,
    /// Minimum similarity threshold for semantic results.
    pub similarity_threshold: f32,
    /// RRF k parameter (default 60).
    pub rrf_k: u32,
    /// Whether to include semantic search.
    pub use_semantic: bool,
    /// Whether to include BM25 search.
    pub use_bm25: bool,
}

impl Default for SearchConfig {
    fn default() -> Self {
        Self {
            top_k: DEFAULT_TOP_K,
            similarity_threshold: DEFAULT_SIMILARITY_THRESHOLD,
            rrf_k: 60,
            use_semantic: true,
            use_bm25: true,
        }
    }
}

/// Default preview length in characters.
pub const DEFAULT_PREVIEW_LEN: usize = 150;

impl SearchResult {
    /// Creates a new search result, looking up chunk metadata from storage.
    ///
    /// Returns `None` if the chunk cannot be found.
    fn from_chunk_id(
        storage: &SqliteStorage,
        chunk_id: i64,
        score: f64,
        semantic_score: Option<f32>,
        bm25_score: Option<f64>,
    ) -> Option<Self> {
        storage
            .get_chunk(chunk_id)
            .ok()
            .flatten()
            .map(|chunk| Self {
                chunk_id,
                buffer_id: chunk.buffer_id,
                index: chunk.index,
                score,
                semantic_score,
                bm25_score,
                content_preview: None,
            })
    }
}

/// Populates content previews for search results.
///
/// # Arguments
///
/// * `storage` - The storage backend.
/// * `results` - Search results to populate.
/// * `preview_len` - Maximum preview length in bytes.
///
/// # Errors
///
/// Returns an error if chunk retrieval fails.
pub fn populate_previews(
    storage: &SqliteStorage,
    results: &mut [SearchResult],
    preview_len: usize,
) -> Result<()> {
    if results.is_empty() {
        return Ok(());
    }

    // Batch-fetch all chunks in a single query instead of N individual lookups.
    let ids: Vec<i64> = results.iter().map(|r| r.chunk_id).collect();
    let chunk_map = storage.get_chunks_by_ids(&ids)?;

    for result in results.iter_mut() {
        if let Some(chunk) = chunk_map.get(&result.chunk_id) {
            let content = &chunk.content;
            let preview = if content.len() <= preview_len {
                content.clone()
            } else {
                // Find a valid UTF-8 boundary
                let end = crate::io::find_char_boundary(content, preview_len);
                let mut preview = content[..end].to_string();
                if end < content.len() {
                    preview.push_str("...");
                }
                preview
            };
            result.content_preview = Some(preview);
        }
    }
    Ok(())
}

impl SearchConfig {
    /// Creates a new search config with default values.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets the top-k limit.
    #[must_use]
    pub const fn with_top_k(mut self, top_k: usize) -> Self {
        self.top_k = top_k;
        self
    }

    /// Sets the similarity threshold.
    #[must_use]
    pub const fn with_threshold(mut self, threshold: f32) -> Self {
        self.similarity_threshold = threshold;
        self
    }

    /// Sets the RRF k parameter.
    #[must_use]
    pub const fn with_rrf_k(mut self, k: u32) -> Self {
        self.rrf_k = k;
        self
    }

    /// Enables or disables semantic search.
    #[must_use]
    pub const fn with_semantic(mut self, enabled: bool) -> Self {
        self.use_semantic = enabled;
        self
    }

    /// Enables or disables BM25 search.
    #[must_use]
    pub const fn with_bm25(mut self, enabled: bool) -> Self {
        self.use_bm25 = enabled;
        self
    }
}

/// Performs hybrid search combining semantic and BM25 results.
///
/// # Arguments
///
/// * `storage` - The storage backend.
/// * `embedder` - The embedding generator.
/// * `query` - The search query text.
/// * `config` - Search configuration.
///
/// # Errors
///
/// Returns an error if search operations fail.
pub fn hybrid_search(
    storage: &SqliteStorage,
    embedder: &dyn Embedder,
    query: &str,
    config: &SearchConfig,
) -> Result<Vec<SearchResult>> {
    hybrid_search_scoped(storage, embedder, query, config, None)
}

/// Performs hybrid search restricted to a single buffer before ranking.
///
/// # Arguments
///
/// * `storage` - The storage backend.
/// * `embedder` - The embedding generator.
/// * `query` - The search query text.
/// * `config` - Search configuration.
/// * `buffer_id` - Buffer ID to restrict before semantic and BM25 ranking.
///
/// # Errors
///
/// Returns an error if search operations fail.
pub fn hybrid_search_in_buffer(
    storage: &SqliteStorage,
    embedder: &dyn Embedder,
    query: &str,
    config: &SearchConfig,
    buffer_id: i64,
) -> Result<Vec<SearchResult>> {
    hybrid_search_scoped(storage, embedder, query, config, Some(buffer_id))
}

fn hybrid_search_scoped(
    storage: &SqliteStorage,
    embedder: &dyn Embedder,
    query: &str,
    config: &SearchConfig,
    buffer_id: Option<i64>,
) -> Result<Vec<SearchResult>> {
    let mut semantic_results: Vec<(i64, f32)> = Vec::new();
    let mut bm25_results: Vec<(i64, f64)> = Vec::new();

    // Semantic search
    if config.use_semantic {
        semantic_results = semantic_search(storage, embedder, query, config, buffer_id)?;
    }

    // BM25 search
    if config.use_bm25 {
        bm25_results = storage.search_fts_in_buffer(query, config.top_k * 2, buffer_id)?;
    }

    // If only one type of search is enabled, return those results directly
    if !config.use_semantic {
        return Ok(bm25_results
            .into_iter()
            .take(config.top_k)
            .filter_map(|(chunk_id, score)| {
                SearchResult::from_chunk_id(storage, chunk_id, score, None, Some(score))
            })
            .collect());
    }

    if !config.use_bm25 {
        return Ok(semantic_results
            .into_iter()
            .take(config.top_k)
            .filter_map(|(chunk_id, score)| {
                SearchResult::from_chunk_id(storage, chunk_id, f64::from(score), Some(score), None)
            })
            .collect());
    }

    // Combine using RRF
    let rrf_config = RrfConfig::new(config.rrf_k);

    // Convert to ranked lists (already sorted by score descending)
    let semantic_ranked: Vec<i64> = semantic_results.iter().map(|(id, _)| *id).collect();
    let bm25_ranked: Vec<i64> = bm25_results.iter().map(|(id, _)| *id).collect();

    let fused = reciprocal_rank_fusion(&[&semantic_ranked, &bm25_ranked], &rrf_config);

    // Build result with original scores
    let semantic_map: std::collections::HashMap<i64, f32> = semantic_results.into_iter().collect();
    let bm25_map: std::collections::HashMap<i64, f64> = bm25_results.into_iter().collect();

    let results: Vec<SearchResult> = fused
        .into_iter()
        .take(config.top_k)
        .filter_map(|(chunk_id, rrf_score)| {
            SearchResult::from_chunk_id(
                storage,
                chunk_id,
                rrf_score,
                semantic_map.get(&chunk_id).copied(),
                bm25_map.get(&chunk_id).copied(),
            )
        })
        .collect();

    Ok(results)
}

/// Performs semantic similarity search.
///
/// Uses cosine similarity between query embedding and stored chunk embeddings.
fn semantic_search(
    storage: &SqliteStorage,
    embedder: &dyn Embedder,
    query: &str,
    config: &SearchConfig,
    buffer_id: Option<i64>,
) -> Result<Vec<(i64, f32)>> {
    use rayon::prelude::*;

    // Generate query embedding
    let query_embedding = embedder.embed(query)?;

    // Get only the embeddings that can participate in this search.
    let all_embeddings = match buffer_id {
        Some(buffer_id) => storage.get_embeddings_in_buffer(buffer_id)?,
        None => storage.get_all_embeddings()?,
    };

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

    // Calculate similarities in parallel (rayon data parallelism)
    let mut similarities: Vec<(i64, f32)> = all_embeddings
        .par_iter()
        .map(|(chunk_id, embedding)| {
            let sim = cosine_similarity(&query_embedding, embedding);
            (*chunk_id, sim)
        })
        .filter(|(_, sim)| *sim >= config.similarity_threshold)
        .collect();

    // Sort by similarity descending
    similarities.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

    // Limit results
    similarities.truncate(config.top_k * 2);

    Ok(similarities)
}

/// Performs semantic-only search.
///
/// # Arguments
///
/// * `storage` - The storage backend.
/// * `embedder` - The embedding generator.
/// * `query` - The search query text.
/// * `top_k` - Maximum number of results.
/// * `threshold` - Minimum similarity threshold.
///
/// # Errors
///
/// Returns an error if search fails.
pub fn search_semantic(
    storage: &SqliteStorage,
    embedder: &dyn Embedder,
    query: &str,
    top_k: usize,
    threshold: f32,
) -> Result<Vec<SearchResult>> {
    let config = SearchConfig::new()
        .with_top_k(top_k)
        .with_threshold(threshold)
        .with_semantic(true)
        .with_bm25(false);

    hybrid_search(storage, embedder, query, &config)
}

/// Performs BM25-only search.
///
/// # Arguments
///
/// * `storage` - The storage backend.
/// * `query` - The search query text.
/// * `top_k` - Maximum number of results.
///
/// # Errors
///
/// Returns an error if search fails.
pub fn search_bm25(
    storage: &SqliteStorage,
    query: &str,
    top_k: usize,
) -> Result<Vec<SearchResult>> {
    let results = storage.search_fts(query, top_k)?;

    Ok(results
        .into_iter()
        .filter_map(|(chunk_id, score)| {
            SearchResult::from_chunk_id(storage, chunk_id, score, None, Some(score))
        })
        .collect())
}

/// Generates and stores embeddings for all chunks in a buffer.
///
/// # Arguments
///
/// * `storage` - The storage backend (mutable for storing embeddings).
/// * `embedder` - The embedding generator.
/// * `buffer_id` - The buffer ID to process.
///
/// # Returns
///
/// The number of chunks embedded.
///
/// # Errors
///
/// Returns an error if embedding generation or storage fails.
pub fn embed_buffer_chunks(
    storage: &mut SqliteStorage,
    embedder: &dyn Embedder,
    buffer_id: i64,
) -> Result<usize> {
    let chunks = storage.get_chunks(buffer_id)?;

    if chunks.is_empty() {
        return Ok(0);
    }

    // Collect chunk texts for batch embedding
    let texts: Vec<&str> = chunks.iter().map(|c| c.content.as_str()).collect();

    // Generate embeddings in batch
    let embeddings = embedder.embed_batch(&texts)?;

    // Prepare batch for storage
    let batch: Vec<(i64, Vec<f32>)> = chunks
        .iter()
        .zip(embeddings)
        .filter_map(|(chunk, embedding)| chunk.id.map(|id| (id, embedding)))
        .collect();

    let count = batch.len();

    // Store embeddings with model name for version tracking
    storage.store_embeddings_batch(&batch, Some(embedder.model_name()))?;

    Ok(count)
}

/// Checks if a buffer has all chunks embedded.
///
/// Uses a single SQL query (`NOT EXISTS`) instead of per-chunk lookups, reducing database
/// round-trips from O(n) to O(1) for a buffer with n chunks.
///
/// # Errors
///
/// Returns an error if the check fails.
pub fn buffer_fully_embedded(storage: &SqliteStorage, buffer_id: i64) -> Result<bool> {
    storage.all_chunks_have_embeddings(buffer_id)
}

/// Checks if existing embeddings were created with a different model.
///
/// Returns `Some(existing_model)` if there's a model mismatch, `None` otherwise.
///
/// # Errors
///
/// Returns an error if the query fails.
pub fn check_model_mismatch(
    storage: &SqliteStorage,
    buffer_id: i64,
    current_model: &str,
) -> Result<Option<String>> {
    let models = storage.get_embedding_models(buffer_id)?;

    // If no embeddings exist, no mismatch
    if models.is_empty() {
        return Ok(None);
    }

    // Check if any existing model differs from the current one
    for model in models {
        if model != current_model {
            return Ok(Some(model));
        }
    }

    Ok(None)
}

/// Information about embedding model versions for a buffer.
#[derive(Debug, Clone)]
pub struct EmbeddingModelInfo {
    /// Model names and their embedding counts.
    pub models: Vec<(Option<String>, i64)>,
    /// Total number of embeddings.
    pub total_embeddings: i64,
    /// Whether there are mixed model versions.
    pub has_mixed_models: bool,
}

/// Gets embedding model information for a buffer.
///
/// # Errors
///
/// Returns an error if the query fails.
pub fn get_embedding_model_info(
    storage: &SqliteStorage,
    buffer_id: i64,
) -> Result<EmbeddingModelInfo> {
    let models = storage.get_embedding_model_counts(buffer_id)?;
    let total_embeddings: i64 = models.iter().map(|(_, count)| count).sum();
    let distinct_models: std::collections::HashSet<_> =
        models.iter().map(|(name, _)| name.as_deref()).collect();
    let has_mixed_models = distinct_models.len() > 1;

    Ok(EmbeddingModelInfo {
        models,
        total_embeddings,
        has_mixed_models,
    })
}

/// Result of an incremental embedding operation.
#[derive(Debug, Clone)]
pub struct IncrementalEmbedResult {
    /// Number of new embeddings created.
    pub embedded_count: usize,
    /// Number of chunks that were skipped (already embedded with correct model).
    pub skipped_count: usize,
    /// Number of embeddings that were replaced (different model).
    pub replaced_count: usize,
    /// Total chunks in the buffer.
    pub total_chunks: usize,
    /// Model name used for embedding.
    pub model_name: String,
}

impl IncrementalEmbedResult {
    /// Returns true if any embeddings were created or updated.
    #[must_use]
    pub const fn had_changes(&self) -> bool {
        self.embedded_count > 0 || self.replaced_count > 0
    }

    /// Returns the percentage of chunks now embedded.
    #[must_use]
    #[allow(clippy::cast_precision_loss)] // Acceptable for percentage calculation
    pub fn completion_percentage(&self) -> f64 {
        if self.total_chunks == 0 {
            100.0
        } else {
            let completed = self.embedded_count + self.skipped_count + self.replaced_count;
            (completed as f64 / self.total_chunks as f64) * 100.0
        }
    }
}

/// Incrementally embeds chunks in a buffer.
///
/// Only embeds chunks that:
/// - Have no embedding, OR
/// - Have an embedding from a different model (if `force_reembed` is true)
///
/// This is more efficient than `embed_buffer_chunks` for large buffers
/// where only a few chunks have changed.
///
/// # Arguments
///
/// * `storage` - The storage backend.
/// * `embedder` - The embedder to use.
/// * `buffer_id` - The buffer to embed.
/// * `force_reembed` - If true, re-embeds chunks with different models.
///
/// # Returns
///
/// An `IncrementalEmbedResult` with statistics about what was done.
///
/// # Errors
///
/// Returns an error if embedding generation or storage fails.
pub fn embed_buffer_chunks_incremental(
    storage: &mut SqliteStorage,
    embedder: &dyn Embedder,
    buffer_id: i64,
    force_reembed: bool,
) -> Result<IncrementalEmbedResult> {
    let current_model = embedder.model_name();
    let stats = storage.get_embedding_stats(buffer_id)?;
    let total_chunks = stats.total_chunks;

    // Determine which chunks need embedding
    let model_to_check = if force_reembed {
        Some(current_model)
    } else {
        None
    };

    let chunk_ids_to_embed = storage.get_chunks_needing_embedding(buffer_id, model_to_check)?;

    if chunk_ids_to_embed.is_empty() {
        return Ok(IncrementalEmbedResult {
            embedded_count: 0,
            skipped_count: total_chunks,
            replaced_count: 0,
            total_chunks,
            model_name: current_model.to_string(),
        });
    }

    // Load the chunks we need to embed
    let all_chunks = storage.get_chunks(buffer_id)?;
    let chunks_to_embed: Vec<_> = all_chunks
        .iter()
        .filter(|c| c.id.is_some_and(|id| chunk_ids_to_embed.contains(&id)))
        .collect();

    // Count how many are replacements (had embeddings before)
    let mut replaced_count = 0;
    for chunk in &chunks_to_embed {
        if let Some(id) = chunk.id
            && storage.has_embedding(id)?
        {
            replaced_count += 1;
        }
    }

    // Generate embeddings
    let texts: Vec<&str> = chunks_to_embed.iter().map(|c| c.content.as_str()).collect();
    let embeddings = embedder.embed_batch(&texts)?;

    // Store embeddings
    let batch: Vec<(i64, Vec<f32>)> = chunks_to_embed
        .iter()
        .zip(embeddings)
        .filter_map(|(chunk, embedding)| chunk.id.map(|id| (id, embedding)))
        .collect();

    let embedded_count = batch.len();
    storage.store_embeddings_batch(&batch, Some(current_model))?;

    let new_embeddings = embedded_count - replaced_count;
    let skipped_count = total_chunks - embedded_count;

    Ok(IncrementalEmbedResult {
        embedded_count: new_embeddings,
        skipped_count,
        replaced_count,
        total_chunks,
        model_name: current_model.to_string(),
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::{Buffer, Chunk};
    use crate::embedding::{DEFAULT_DIMENSIONS, Embedder, FallbackEmbedder};
    use crate::storage::Storage;

    fn setup_storage() -> SqliteStorage {
        let mut storage = SqliteStorage::in_memory().unwrap();
        storage.init().unwrap();
        storage
    }

    fn setup_storage_with_chunks() -> SqliteStorage {
        let mut storage = setup_storage();

        // Create a buffer
        let buffer = Buffer::from_named(
            "test.txt".to_string(),
            "Test content for searching".to_string(),
        );
        let buffer_id = storage.add_buffer(&buffer).unwrap();

        // Create chunks with different content
        let chunks = vec![
            Chunk::new(
                buffer_id,
                "The quick brown fox jumps over the lazy dog".to_string(),
                0..44,
                0,
            ),
            Chunk::new(
                buffer_id,
                "Machine learning is a subset of artificial intelligence".to_string(),
                44..100,
                1,
            ),
            Chunk::new(
                buffer_id,
                "Rust is a systems programming language".to_string(),
                100..139,
                2,
            ),
        ];

        storage.add_chunks(buffer_id, &chunks).unwrap();

        storage
    }

    fn add_single_chunk_buffer(storage: &mut SqliteStorage, content: &str) -> (i64, i64) {
        let buffer_id = storage
            .add_buffer(&Buffer::from_content(content.to_string()))
            .unwrap();
        storage
            .add_chunks(
                buffer_id,
                &[Chunk::new(
                    buffer_id,
                    content.to_string(),
                    0..content.len(),
                    0,
                )],
            )
            .unwrap();
        let chunk_id = storage.get_chunks(buffer_id).unwrap()[0].id.unwrap();
        (buffer_id, chunk_id)
    }

    #[test]
    fn test_search_config_default() {
        let config = SearchConfig::default();
        assert_eq!(config.top_k, DEFAULT_TOP_K);
        assert!((config.similarity_threshold - DEFAULT_SIMILARITY_THRESHOLD).abs() < f32::EPSILON);
        assert_eq!(config.rrf_k, 60);
        assert!(config.use_semantic);
        assert!(config.use_bm25);
    }

    #[test]
    fn test_search_config_builder() {
        let config = SearchConfig::new()
            .with_top_k(20)
            .with_threshold(0.5)
            .with_rrf_k(30)
            .with_semantic(false)
            .with_bm25(true);

        assert_eq!(config.top_k, 20);
        assert!((config.similarity_threshold - 0.5).abs() < f32::EPSILON);
        assert_eq!(config.rrf_k, 30);
        assert!(!config.use_semantic);
        assert!(config.use_bm25);
    }

    #[test]
    fn test_search_bm25() {
        let storage = setup_storage_with_chunks();

        // Search for "fox" - should find the first chunk
        let results = search_bm25(&storage, "fox", 10).unwrap();
        assert!(!results.is_empty());
        assert!(results[0].bm25_score.is_some());
        assert!(results[0].semantic_score.is_none());
    }

    #[test]
    fn test_search_bm25_no_results() {
        let storage = setup_storage_with_chunks();

        // Search for something not in the content
        let results = search_bm25(&storage, "xyz123nonexistent", 10).unwrap();
        assert!(results.is_empty());
    }

    #[test]
    fn test_embed_buffer_chunks() {
        let mut storage = setup_storage_with_chunks();
        let embedder = FallbackEmbedder::new(DEFAULT_DIMENSIONS);

        // Embed chunks for buffer 1
        let count = embed_buffer_chunks(&mut storage, &embedder, 1).unwrap();
        assert_eq!(count, 3); // We created 3 chunks
    }

    #[test]
    fn test_embed_buffer_chunks_empty() {
        let mut storage = setup_storage();
        let embedder = FallbackEmbedder::new(DEFAULT_DIMENSIONS);

        // Create buffer with no chunks
        let buffer = Buffer::from_named("empty.txt".to_string(), String::new());
        let buffer_id = storage.add_buffer(&buffer).unwrap();

        let count = embed_buffer_chunks(&mut storage, &embedder, buffer_id).unwrap();
        assert_eq!(count, 0);
    }

    #[test]
    fn test_buffer_fully_embedded_empty() {
        let mut storage = setup_storage();

        // Create buffer with no chunks
        let buffer = Buffer::from_named("empty.txt".to_string(), String::new());
        let buffer_id = storage.add_buffer(&buffer).unwrap();

        // Empty buffer should be "fully embedded"
        let result = buffer_fully_embedded(&storage, buffer_id).unwrap();
        assert!(result);
    }

    #[test]
    fn test_buffer_fully_embedded_with_embeddings() {
        let mut storage = setup_storage_with_chunks();
        let embedder = FallbackEmbedder::new(DEFAULT_DIMENSIONS);

        // Before embedding
        let result = buffer_fully_embedded(&storage, 1).unwrap();
        assert!(!result);

        // Embed all chunks
        embed_buffer_chunks(&mut storage, &embedder, 1).unwrap();

        // After embedding
        let result = buffer_fully_embedded(&storage, 1).unwrap();
        assert!(result);
    }

    #[test]
    fn test_hybrid_search_bm25_only() {
        let storage = setup_storage_with_chunks();
        let embedder = FallbackEmbedder::new(DEFAULT_DIMENSIONS);

        let config = SearchConfig::new().with_semantic(false).with_bm25(true);

        let results = hybrid_search(&storage, &embedder, "programming", &config).unwrap();
        // Should find "Rust is a systems programming language"
        assert!(!results.is_empty());
        assert!(results[0].bm25_score.is_some());
        assert!(results[0].semantic_score.is_none());
    }

    #[test]
    fn test_buffer_filter_applies_before_ranking() {
        let mut storage = setup_storage();
        let (target_buffer, target_chunk) = add_single_chunk_buffer(&mut storage, "needle");
        let (_, distractor_chunk) =
            add_single_chunk_buffer(&mut storage, "needle needle needle needle");
        let embedder = FallbackEmbedder::new(2);
        let query_embedding = embedder.embed("needle").unwrap();
        // The target embedding is orthogonal to the query while the distractor
        // matches, so this fails if buffer filtering happens after ranking.
        storage
            .store_embedding(
                target_chunk,
                &[-query_embedding[1], query_embedding[0]],
                Some("query-test"),
            )
            .unwrap();
        storage
            .store_embedding(distractor_chunk, &query_embedding, Some("query-test"))
            .unwrap();

        let bm25_config = SearchConfig::new()
            .with_top_k(1)
            .with_semantic(false)
            .with_bm25(true);
        let bm25_results =
            hybrid_search_in_buffer(&storage, &embedder, "needle", &bm25_config, target_buffer)
                .unwrap();
        assert_eq!(bm25_results[0].chunk_id, target_chunk);

        let semantic_config = SearchConfig::new()
            .with_top_k(1)
            .with_threshold(-1.0)
            .with_semantic(true)
            .with_bm25(false);
        let semantic_results = hybrid_search_in_buffer(
            &storage,
            &embedder,
            "needle",
            &semantic_config,
            target_buffer,
        )
        .unwrap();
        assert_eq!(semantic_results[0].chunk_id, target_chunk);
    }

    #[test]
    fn test_hybrid_search_semantic_only() {
        let mut storage = setup_storage_with_chunks();
        let embedder = FallbackEmbedder::new(DEFAULT_DIMENSIONS);

        // First embed the chunks
        embed_buffer_chunks(&mut storage, &embedder, 1).unwrap();

        let config = SearchConfig::new()
            .with_semantic(true)
            .with_bm25(false)
            .with_threshold(0.0); // Low threshold for fallback embedder

        let results = hybrid_search(&storage, &embedder, "programming language", &config).unwrap();
        assert!(!results.is_empty());
        assert!(results[0].semantic_score.is_some());
        assert!(results[0].bm25_score.is_none());
    }

    #[test]
    fn test_hybrid_search_both() {
        let mut storage = setup_storage_with_chunks();
        let embedder = FallbackEmbedder::new(DEFAULT_DIMENSIONS);

        // First embed the chunks
        embed_buffer_chunks(&mut storage, &embedder, 1).unwrap();

        let config = SearchConfig::new()
            .with_semantic(true)
            .with_bm25(true)
            .with_threshold(0.0); // Low threshold for fallback embedder

        let results = hybrid_search(&storage, &embedder, "programming", &config).unwrap();
        assert!(!results.is_empty());
    }

    #[test]
    fn test_search_semantic() {
        let mut storage = setup_storage_with_chunks();
        let embedder = FallbackEmbedder::new(DEFAULT_DIMENSIONS);

        // First embed the chunks
        embed_buffer_chunks(&mut storage, &embedder, 1).unwrap();

        let results = search_semantic(&storage, &embedder, "test query", 10, 0.0).unwrap();
        // Should return results with semantic scores only
        for result in &results {
            assert!(result.semantic_score.is_some());
            assert!(result.bm25_score.is_none());
        }
    }

    #[test]
    fn test_search_semantic_empty_embeddings() {
        let storage = setup_storage_with_chunks();
        let embedder = FallbackEmbedder::new(DEFAULT_DIMENSIONS);

        // Don't embed chunks - search should return empty
        let results = search_semantic(&storage, &embedder, "test query", 10, 0.5).unwrap();
        assert!(results.is_empty());
    }

    #[test]
    fn test_incremental_embed_new_chunks() {
        let mut storage = setup_storage_with_chunks();
        let embedder = FallbackEmbedder::new(DEFAULT_DIMENSIONS);

        // First incremental embed - should embed all 3 chunks
        let result = embed_buffer_chunks_incremental(&mut storage, &embedder, 1, false).unwrap();
        assert_eq!(result.embedded_count, 3);
        assert_eq!(result.skipped_count, 0);
        assert_eq!(result.replaced_count, 0);
        assert_eq!(result.total_chunks, 3);
        assert!(result.had_changes());

        // Second incremental embed - should skip all (already embedded)
        let result2 = embed_buffer_chunks_incremental(&mut storage, &embedder, 1, false).unwrap();
        assert_eq!(result2.embedded_count, 0);
        assert_eq!(result2.skipped_count, 3);
        assert_eq!(result2.replaced_count, 0);
        assert!(!result2.had_changes());
    }

    #[test]
    fn test_incremental_embed_force_reembed() {
        let mut storage = setup_storage_with_chunks();
        let embedder = FallbackEmbedder::new(DEFAULT_DIMENSIONS);

        // First embed normally
        embed_buffer_chunks_incremental(&mut storage, &embedder, 1, false).unwrap();

        // Force re-embed - should replace all 3
        let result = embed_buffer_chunks_incremental(&mut storage, &embedder, 1, true).unwrap();
        // All chunks already have correct model, so no changes needed even with force
        // (force only affects different-model embeddings)
        assert_eq!(result.skipped_count, 3);
        assert!(!result.had_changes());
    }

    #[test]
    fn test_parallel_semantic_search() {
        let mut storage = setup_storage_with_chunks();
        let embedder = FallbackEmbedder::new(DEFAULT_DIMENSIONS);

        // Embed all chunks so semantic_search has data to work with
        embed_buffer_chunks(&mut storage, &embedder, 1).unwrap();

        // Use a zero threshold to capture all results from the fallback embedder
        let config = SearchConfig::new().with_top_k(10).with_threshold(0.0);

        let results = semantic_search(&storage, &embedder, "test query", &config, None).unwrap();

        // Should return results (we embedded 3 chunks)
        assert!(!results.is_empty());

        // Verify results are sorted by descending similarity
        for window in results.windows(2) {
            assert!(
                window[0].1 >= window[1].1,
                "Results should be sorted by descending similarity: {} >= {}",
                window[0].1,
                window[1].1,
            );
        }

        // Verify threshold filtering: with a very high threshold, results should be excluded
        let strict_config = SearchConfig::new().with_top_k(10).with_threshold(0.99);
        let strict_results =
            semantic_search(&storage, &embedder, "test query", &strict_config, None).unwrap();

        // Strict threshold should return fewer (or no) results
        assert!(
            strict_results.len() <= results.len(),
            "Strict threshold should not return more results than lenient threshold",
        );
    }

    #[test]
    fn test_incremental_embed_result_completion() {
        let result = IncrementalEmbedResult {
            embedded_count: 2,
            skipped_count: 3,
            replaced_count: 0,
            total_chunks: 5,
            model_name: "test".to_string(),
        };
        assert!(result.had_changes());
        assert!((result.completion_percentage() - 100.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_completion_percentage_zero_chunks() {
        let result = IncrementalEmbedResult {
            embedded_count: 0,
            skipped_count: 0,
            replaced_count: 0,
            total_chunks: 0,
            model_name: "test".to_string(),
        };
        // Zero-chunk buffer is considered 100% complete
        assert!((result.completion_percentage() - 100.0).abs() < f64::EPSILON);
        assert!(!result.had_changes());
    }

    #[test]
    fn test_completion_percentage_partial() {
        let result = IncrementalEmbedResult {
            embedded_count: 1,
            skipped_count: 1,
            replaced_count: 0,
            total_chunks: 4,
            model_name: "test".to_string(),
        };
        assert!((result.completion_percentage() - 50.0).abs() < f64::EPSILON);
        assert!(result.had_changes());
    }

    #[test]
    fn test_had_changes_replaced_only() {
        let result = IncrementalEmbedResult {
            embedded_count: 0,
            skipped_count: 2,
            replaced_count: 1,
            total_chunks: 3,
            model_name: "test".to_string(),
        };
        assert!(result.had_changes());
    }

    #[test]
    fn test_check_model_mismatch_no_embeddings() {
        let mut storage = setup_storage();
        let buffer = crate::core::Buffer::from_named("x.txt".to_string(), "hi".to_string());
        let buffer_id = storage.add_buffer(&buffer).unwrap();

        // No embeddings → no mismatch
        let result = check_model_mismatch(&storage, buffer_id, "model-a").unwrap();
        assert!(result.is_none());
    }

    #[test]
    fn test_check_model_mismatch_same_model() {
        let mut storage = setup_storage_with_chunks();
        let embedder = FallbackEmbedder::new(DEFAULT_DIMENSIONS);

        embed_buffer_chunks(&mut storage, &embedder, 1).unwrap();

        // Same model → no mismatch
        let result = check_model_mismatch(&storage, 1, embedder.model_name()).unwrap();
        assert!(result.is_none());
    }

    #[test]
    fn test_check_model_mismatch_different_model() {
        let mut storage = setup_storage_with_chunks();
        let embedder = FallbackEmbedder::new(DEFAULT_DIMENSIONS);

        embed_buffer_chunks(&mut storage, &embedder, 1).unwrap();

        // Ask with a different model name → returns the existing model
        let result = check_model_mismatch(&storage, 1, "some-other-model").unwrap();
        assert!(result.is_some());
        assert_eq!(result.unwrap(), embedder.model_name());
    }

    #[test]
    fn test_get_embedding_model_info_no_embeddings() {
        let mut storage = setup_storage();
        let buffer = crate::core::Buffer::from_named("x.txt".to_string(), "hi".to_string());
        let buffer_id = storage.add_buffer(&buffer).unwrap();

        let info = get_embedding_model_info(&storage, buffer_id).unwrap();
        assert_eq!(info.total_embeddings, 0);
        assert!(info.models.is_empty());
        assert!(!info.has_mixed_models);
    }

    #[test]
    fn test_get_embedding_model_info_single_model() {
        let mut storage = setup_storage_with_chunks();
        let embedder = FallbackEmbedder::new(DEFAULT_DIMENSIONS);

        embed_buffer_chunks(&mut storage, &embedder, 1).unwrap();

        let info = get_embedding_model_info(&storage, 1).unwrap();
        assert_eq!(info.total_embeddings, 3);
        assert!(!info.has_mixed_models);
    }

    #[test]
    fn test_populate_previews_short_content() {
        let storage = setup_storage_with_chunks();

        // Get a result with a known chunk_id
        let results_raw = search_bm25(&storage, "fox", 1).unwrap();
        assert!(!results_raw.is_empty());

        let mut results = results_raw;
        populate_previews(&storage, &mut results, 200).unwrap();

        // Content is shorter than preview_len → no ellipsis
        let preview = results[0].content_preview.as_ref().unwrap();
        assert!(!preview.ends_with("..."));
        assert!(!preview.is_empty());
    }

    #[test]
    fn test_populate_previews_truncates_long_content() {
        let mut storage = setup_storage();
        let buffer = crate::core::Buffer::from_named("long.txt".to_string(), String::new());
        let buffer_id = storage.add_buffer(&buffer).unwrap();

        // Create a chunk with long content
        let long_content = "word ".repeat(100); // 500 chars
        let chunk =
            crate::core::Chunk::new(buffer_id, long_content.clone(), 0..long_content.len(), 0);
        storage.add_chunks(buffer_id, &[chunk]).unwrap();

        let chunks = storage.get_chunks(buffer_id).unwrap();
        let chunk_id = chunks[0].id.unwrap();

        // Build a synthetic SearchResult pointing at this chunk
        let mut results = vec![SearchResult {
            chunk_id,
            buffer_id,
            index: 0,
            score: 1.0,
            semantic_score: None,
            bm25_score: None,
            content_preview: None,
        }];

        populate_previews(&storage, &mut results, 20).unwrap();

        let preview = results[0].content_preview.as_ref().unwrap();
        assert!(
            preview.ends_with("..."),
            "Expected ellipsis, got: {preview}"
        );
        // Should be no longer than preview_len + "..."
        assert!(preview.len() <= 23);
    }

    #[test]
    fn test_populate_previews_utf8_boundary() {
        let mut storage = setup_storage();
        let buffer = crate::core::Buffer::from_named("utf8.txt".to_string(), String::new());
        let buffer_id = storage.add_buffer(&buffer).unwrap();

        // Content where a naive byte truncation would split a multi-byte char
        // '日' is 3 bytes; place it so it straddles the preview boundary
        let content = "hello \u{65E5}\u{672C}\u{8A9E}"; // "hello 日本語"
        let chunk = crate::core::Chunk::new(buffer_id, content.to_string(), 0..content.len(), 0);
        storage.add_chunks(buffer_id, &[chunk]).unwrap();

        let chunks = storage.get_chunks(buffer_id).unwrap();
        let chunk_id = chunks[0].id.unwrap();

        let mut results = vec![SearchResult {
            chunk_id,
            buffer_id,
            index: 0,
            score: 1.0,
            semantic_score: None,
            bm25_score: None,
            content_preview: None,
        }];

        // preview_len=7 bytes falls inside '日' (which occupies bytes 6-8),
        // so the implementation must back up to the nearest valid UTF-8 boundary.
        populate_previews(&storage, &mut results, 7).unwrap();

        let preview = results[0].content_preview.as_ref().unwrap();
        // Must end with ellipsis because the content is longer than preview_len
        assert!(preview.ends_with("..."), "Expected ellipsis in: {preview}");
        // The prefix before "..." must be exactly "hello " — truncated at the
        // char boundary before '日', not mid-way through its bytes.
        let body = preview.trim_end_matches("...");
        assert_eq!(
            body, "hello ",
            "Expected truncation at char boundary, got: {body:?}"
        );
    }
}