//! HNSW index construction

use crate::embedder::Embedding;
use crate::index::DistanceMetric;

use super::{HnswError, HnswGraph, HnswIndex, HnswInner, MAX_LAYER};

impl HnswIndex {
    /// Build a new HNSW index from embeddings (single-pass).
    ///
    /// # When to use `build_with_dim` vs `build_batched_with_dim`
    ///
    /// - **`build_with_dim`**: Use when all embeddings fit comfortably in memory
    ///   (<50k chunks, ~150MB for 50k x 768 x 4 bytes). Slightly higher graph
    ///   quality since all vectors are available during construction.
    ///
    /// - **`build_batched_with_dim`**: Use for large indexes (>50k chunks) or
    ///   memory-constrained environments. Streams embeddings in batches to avoid
    ///   OOM. Graph quality is marginally lower but negligible for practical
    ///   search accuracy.
    ///
    /// **Warning:** This loads all embeddings into memory at once.
    /// For large indexes (>50k chunks), prefer `build_batched_with_dim()` to
    /// avoid OOM.
    ///
    /// # Production routing
    ///
    /// `build_hnsw_index()` in `cli/commands/index.rs` unconditionally uses
    /// `build_batched_with_dim()` with 10k-row batches for all index sizes.
    /// This method is only used in tests.
    ///
    /// # Arguments
    /// * `embeddings` - Vector of (chunk_id, embedding) pairs
    /// * `dim` - Expected embedding dimension
    ///
    /// The distance metric is resolved from `CQS_DISTANCE_METRIC` (default
    /// cosine) — see [`DistanceMetric::resolve`]. Use
    /// [`Self::build_with_dim_and_metric`] to pin it explicitly.
    pub fn build_with_dim(
        embeddings: Vec<(String, Embedding)>,
        dim: usize,
    ) -> Result<Self, HnswError> {
        let metric = DistanceMetric::resolve().map_err(HnswError::Build)?;
        Self::build_with_dim_and_metric(embeddings, dim, metric)
    }

    /// [`Self::build_with_dim`] with an explicit [`DistanceMetric`] instead
    /// of the `CQS_DISTANCE_METRIC` env resolution.
    pub fn build_with_dim_and_metric(
        embeddings: Vec<(String, Embedding)>,
        dim: usize,
        metric: DistanceMetric,
    ) -> Result<Self, HnswError> {
        let _span = tracing::debug_span!("hnsw_build", metric = %metric).entered();
        if dim == 0 {
            return Err(HnswError::Build("Embedding dimension must be > 0".into()));
        }
        if embeddings.is_empty() {
            // Create empty index — fall back to the small-tier default.
            let hnsw = HnswGraph::new(
                metric,
                super::max_nb_connection_for(0),
                1,
                MAX_LAYER,
                super::ef_construction_for(0),
            );
            return Ok(Self {
                inner: HnswInner::Owned(hnsw),
                id_map: Vec::new(),
                ef_search: super::ef_search_for(0),
                dim,
                _lock_file: None,
            });
        }

        let (id_map, data, nb_elem) = super::prepare_index_data(embeddings, dim)?;

        tracing::info!(count = nb_elem, metric = %metric, "Building HNSW index");

        // Corpus-size-aware tier defaults. Env overrides still win — see
        // `hnsw_tier_defaults` for the table.
        let mut hnsw = HnswGraph::new(
            metric,
            super::max_nb_connection_for(nb_elem),
            nb_elem,
            MAX_LAYER,
            super::ef_construction_for(nb_elem),
        );

        // Test-only path: allocates the full Vec<Vec<f32>> double-buffer here.
        // Production code uses `build_batched_with_dim` to avoid this peak
        // allocation.
        // Reconstruct Vec<f32> chunks from flat buffer for hnsw_rs API
        let chunks: Vec<Vec<f32>> = data.chunks_exact(dim).map(|c| c.to_vec()).collect();
        let data_for_insert: Vec<(&Vec<f32>, usize)> =
            chunks.iter().enumerate().map(|(i, v)| (v, i)).collect();

        // Parallel insert for performance
        hnsw.parallel_insert_data(&data_for_insert);

        tracing::info!("HNSW index built successfully");

        Ok(Self {
            inner: HnswInner::Owned(hnsw),
            ef_search: super::ef_search_for(id_map.len()),
            id_map,
            dim,
            _lock_file: None,
        })
    }

    /// Build HNSW index incrementally from batches (memory-efficient).
    ///
    /// Processes embeddings in batches to avoid loading everything into RAM.
    /// Each batch is inserted via `parallel_insert`, building the graph incrementally.
    ///
    /// Memory usage: O(batch_size) instead of O(total_embeddings).
    /// Trade-off: Slightly lower graph quality vs. single-pass build, but
    /// negligible for practical search accuracy.
    ///
    /// # Arguments
    /// * `batches` - Iterator yielding `Result<Vec<(id, embedding)>>` batches
    /// * `estimated_total` - Hint for HNSW capacity (can be approximate)
    /// * `dim` - Expected embedding dimension
    ///
    /// # Example
    ///
    /// `text` rather than `ignore` because the snippet references `store`
    /// without setup; `cargo test -- --include-ignored` (the `ci-slow.yml`
    /// shape) compiles `ignore`-tagged doctests and would surface
    /// `cannot find value 'store' in this scope`.
    ///
    /// ```text
    /// let index = HnswIndex::build_batched_with_dim(
    ///     store.embedding_batches(10_000),
    ///     store.chunk_count()?,
    ///     768,
    /// )?;
    /// ```
    pub fn build_batched_with_dim<I, E>(
        batches: I,
        estimated_total: usize,
        dim: usize,
    ) -> Result<Self, HnswError>
    where
        I: Iterator<Item = Result<Vec<(String, Embedding)>, E>>,
        E: std::fmt::Display,
    {
        let metric = DistanceMetric::resolve().map_err(HnswError::Build)?;
        Self::build_batched_with_dim_and_metric(batches, estimated_total, dim, metric)
    }

    /// [`Self::build_batched_with_dim`] with an explicit [`DistanceMetric`]
    /// instead of the `CQS_DISTANCE_METRIC` env resolution.
    pub fn build_batched_with_dim_and_metric<I, E>(
        batches: I,
        estimated_total: usize,
        dim: usize,
        metric: DistanceMetric,
    ) -> Result<Self, HnswError>
    where
        I: Iterator<Item = Result<Vec<(String, Embedding)>, E>>,
        E: std::fmt::Display,
    {
        let _span =
            tracing::debug_span!("hnsw_build_batched", estimated_total, metric = %metric).entered();
        if dim == 0 {
            return Err(HnswError::Build("Embedding dimension must be > 0".into()));
        }
        let capacity = estimated_total.max(1);
        tracing::info!(
            metric = %metric,
            "Building HNSW index incrementally (estimated {} vectors)",
            capacity
        );

        // Tier defaults read from `capacity`, the estimated vector count
        // for this build.
        let mut hnsw = HnswGraph::new(
            metric,
            super::max_nb_connection_for(capacity),
            capacity,
            MAX_LAYER,
            super::ef_construction_for(capacity),
        );

        let mut id_map: Vec<Box<str>> = Vec::with_capacity(capacity);
        let mut total_inserted = 0usize;
        let mut batch_num = 0usize;

        for batch_result in batches {
            let batch = batch_result
                .map_err(|e| HnswError::Internal(format!("Batch fetch failed: {}", e)))?;

            if batch.is_empty() {
                continue;
            }

            // Validate dimensions and build insertion data in a single pass.
            // Use a separate insertion counter (not the loop index) because
            // zero-vector skips would desync base_idx+i from id_map positions.
            let mut data_for_insert: Vec<(&Vec<f32>, usize)> = Vec::with_capacity(batch.len());

            for (chunk_id, embedding) in batch.iter() {
                if embedding.len() != dim {
                    return Err(HnswError::DimensionMismatch {
                        expected: dim,
                        actual: embedding.len(),
                    });
                }
                // Skip zero-vector embeddings — they produce NaN cosine
                // distances. Short-circuit on the first non-zero element
                // instead of computing the full L2 norm. Note `NaN != 0.0`
                // is true, so this check alone does not catch NaN vectors.
                if !embedding.as_vec().iter().any(|x| *x != 0.0) {
                    tracing::warn!(chunk_id = %chunk_id, "Skipping zero-vector embedding");
                    continue;
                }
                // Skip embeddings with NaN/Inf components — non-finite
                // distances trip assertions inside hnsw_rs's parallel insert
                // workers and would panic the whole build.
                if embedding.as_vec().iter().any(|x| !x.is_finite()) {
                    tracing::warn!(chunk_id = %chunk_id, "Skipping non-finite embedding");
                    continue;
                }
                let insert_idx = id_map.len();
                tracing::trace!("Adding {} to HNSW index at {}", chunk_id, insert_idx);
                // Clone bytes into Box<str> (no cap field).
                id_map.push(Box::from(chunk_id.as_str()));
                data_for_insert.push((embedding.as_vec(), insert_idx));
            }

            // Insert this batch (hnsw_rs supports consecutive parallel_insert calls)
            hnsw.parallel_insert_data(&data_for_insert);

            total_inserted += data_for_insert.len();
            batch_num += 1;
            tracing::debug!(
                batch = batch_num,
                vectors_so_far = total_inserted,
                "HNSW batch inserted"
            );
            let progress_pct = (total_inserted * 100).checked_div(capacity).unwrap_or(100);
            // debug! not info!. A 60k-chunk build emits 50+ batches; at info
            // level they drown out meaningful log messages with no
            // user-actionable signal per batch. The final "HNSW index built"
            // message (info!) is sufficient for progress tracking.
            tracing::debug!(
                "HNSW build progress: {} / ~{} vectors ({}%)",
                total_inserted,
                capacity,
                progress_pct
            );
        }

        if id_map.is_empty() {
            tracing::info!("HNSW index built (empty)");
            return Ok(Self {
                inner: HnswInner::Owned(HnswGraph::new(
                    metric,
                    super::max_nb_connection_for(0),
                    1,
                    MAX_LAYER,
                    super::ef_construction_for(0),
                )),
                id_map: Vec::new(),
                ef_search: super::ef_search_for(0),
                dim,
                _lock_file: None,
            });
        }

        tracing::info!(count = id_map.len(), "HNSW index built");

        Ok(Self {
            inner: HnswInner::Owned(hnsw),
            ef_search: super::ef_search_for(id_map.len()),
            id_map,
            dim,
            _lock_file: None,
        })
    }
}

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

    use crate::hnsw::assert_self_match_reachable;
    use crate::hnsw::make_test_embedding as make_embedding;

    #[test]
    fn test_build_and_search() {
        // Recall-intent test: a correctly-built index returns the exact-match
        // vector in its top-k. `assert_self_match_reachable` retries the build
        // (and takes HNSW_ENV_LOCK around it) to absorb the rare degenerate
        // concurrent-build graph — see the helper docs in hnsw/mod.rs.
        let build = || {
            let embeddings = vec![
                ("chunk1".to_string(), make_embedding(1)),
                ("chunk2".to_string(), make_embedding(2)),
                ("chunk3".to_string(), make_embedding(3)),
            ];
            let index = HnswIndex::build_with_dim(embeddings, crate::EMBEDDING_DIM).unwrap();
            assert_eq!(index.len(), 3);
            index
        };

        let self_score = assert_self_match_reachable(build, &make_embedding(1), "chunk1", 3);
        // Self-match score must be high (cosine distance ~0).
        assert!(
            self_score > 0.9,
            "self-match score should be high, got {self_score}"
        );
    }

    #[test]
    fn test_empty_index() {
        let index = HnswIndex::build_with_dim(vec![], crate::EMBEDDING_DIM).unwrap();
        assert!(index.is_empty());

        let query = make_embedding(1);
        let results = index.search(&query, 5);
        assert!(results.is_empty());
    }

    #[test]
    fn test_build_batched() {
        // Simulate streaming batches like Store::embedding_batches would
        // provide. Seeds spaced by 1000 so the sin-based embeddings are
        // genuinely distinct in cosine space — see the rationale on
        // `test_build_batched_vs_regular_equivalence`.
        let all_embeddings: Vec<(String, Embedding)> = (1..=25)
            .map(|i| (format!("chunk{}", i), make_embedding(i * 1000)))
            .collect();

        // Split into batches of 10 (simulating LIMIT/OFFSET pagination)
        let batches: Vec<Result<Vec<(String, Embedding)>, std::convert::Infallible>> =
            all_embeddings
                .chunks(10)
                .map(|chunk| Ok(chunk.to_vec()))
                .collect();

        let index =
            HnswIndex::build_batched_with_dim(batches.into_iter(), 25, crate::EMBEDDING_DIM)
                .unwrap();
        assert_eq!(index.len(), 25);

        // Wiring check: query the seed directly — nearest neighbour must
        // be itself. With spaced seeds the small-tier graph reliably
        // returns the queried point in top-K=N.
        let query = make_embedding(1000);
        let results = index.search(&query, 25);
        assert!(!results.is_empty());
        assert!(
            results.iter().any(|r| r.id == "chunk1"),
            "chunk1 missing from top-25 of N=25; got {} results: {:?}",
            results.len(),
            results.iter().map(|r| &r.id).collect::<Vec<_>>()
        );
    }

    #[test]
    fn test_build_batched_empty() {
        let batches: Vec<Result<Vec<(String, Embedding)>, std::convert::Infallible>> = vec![];
        let index = HnswIndex::build_batched_with_dim(batches.into_iter(), 0, crate::EMBEDDING_DIM)
            .unwrap();
        assert!(index.is_empty());
    }

    #[test]
    fn test_build_batched_vs_regular_equivalence() {
        // Build same index both ways, verify similar search results.
        //
        // Seeds spaced by 1000 so the sin-based embeddings are genuinely
        // distinct in cosine space. Adjacent seeds (1..=20) produce
        // near-collinear vectors that flake on the small-tier HNSW defaults
        // (M=16, ef_c=100, ef_s=50 for <5k chunks): sin(0.1) ≈ sin(0.2),
        // and HNSW's approximate graph drops one of them out of recall
        // non-deterministically. Spaced seeds give sin(100), sin(200), ...
        // — uniformly distributed in [-1, 1] — so recall@N is deterministic.
        let item_id = |i: u32| format!("item{}", i);
        let embeddings: Vec<(String, Embedding)> = (1..=20)
            .map(|i| (item_id(i), make_embedding(i * 1000)))
            .collect();

        let regular = HnswIndex::build_with_dim(embeddings.clone(), crate::EMBEDDING_DIM).unwrap();

        let batches: Vec<Result<Vec<(String, Embedding)>, std::convert::Infallible>> = embeddings
            .chunks(7) // Odd batch size to test edge cases
            .map(|chunk| Ok(chunk.to_vec()))
            .collect();
        let batched =
            HnswIndex::build_batched_with_dim(batches.into_iter(), 20, crate::EMBEDDING_DIM)
                .unwrap();

        assert_eq!(regular.len(), batched.len());

        // Both should find item10 — query embedding matches its seed
        // exactly, so the nearest neighbour is itself.
        let query = make_embedding(10 * 1000);
        let regular_results = regular.search(&query, 20);
        let batched_results = batched.search(&query, 20);

        let regular_found = regular_results.iter().any(|r| r.id == "item10");
        let batched_found = batched_results.iter().any(|r| r.id == "item10");
        assert!(regular_found, "Regular build should find item10 in top 20");
        assert!(batched_found, "Batched build should find item10 in top 20");
    }

    // ===== multi-model dim-threading (HNSW build) =====

    /// Create a deterministic normalized embedding of arbitrary dimension.
    fn make_embedding_dim(seed: u32, dim: usize) -> Embedding {
        let mut v = vec![0.0f32; dim];
        for (i, val) in v.iter_mut().enumerate() {
            *val = ((seed as f32 * 0.1) + (i as f32 * 0.001)).sin();
        }
        let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        if norm > 0.0 {
            for val in &mut v {
                *val /= norm;
            }
        }
        Embedding::new(v)
    }

    #[test]
    fn tc31_build_batched_with_dim_1024() {
        // Build HNSW index with 1024-dim embeddings via build_batched_with_dim.
        let build = || {
            let all_embeddings: Vec<(String, Embedding)> = (1..=10)
                .map(|i| (format!("chunk{}", i), make_embedding_dim(i, 1024)))
                .collect();
            let batches: Vec<Result<Vec<(String, Embedding)>, std::convert::Infallible>> =
                all_embeddings
                    .chunks(5)
                    .map(|chunk| Ok(chunk.to_vec()))
                    .collect();
            let index = HnswIndex::build_batched_with_dim(batches.into_iter(), 10, 1024).unwrap();
            assert_eq!(index.len(), 10, "HNSW index should have 10 vectors");
            assert_eq!(index.dim, 1024, "HNSW index dim should be 1024");
            index
        };
        // Recall check with bounded build retries.
        assert_self_match_reachable(build, &make_embedding_dim(1, 1024), "chunk1", 3);
    }

    #[test]
    fn tc31_build_with_dim_1024() {
        // Single-pass build with 1024-dim.
        let build = || {
            let embeddings: Vec<(String, Embedding)> = (1..=5)
                .map(|i| (format!("item{}", i), make_embedding_dim(i, 1024)))
                .collect();
            let index = HnswIndex::build_with_dim(embeddings, 1024).unwrap();
            assert_eq!(index.len(), 5);
            assert_eq!(index.dim, 1024);
            index
        };
        // Recall check with bounded build retries.
        assert_self_match_reachable(build, &make_embedding_dim(3, 1024), "item3", 5);
    }

    #[test]
    fn tc31_build_batched_dim_mismatch_rejected() {
        // Feeding 128-dim embeddings to a 1024-dim build should fail.
        let bad_embeddings: Vec<(String, Embedding)> = (1..=3)
            .map(|i| {
                let mut v = vec![0.0f32; 128]; // intentionally wrong dim
                v[0] = i as f32;
                (format!("chunk{}", i), Embedding::new(v))
            })
            .collect();

        let batches: Vec<Result<Vec<(String, Embedding)>, std::convert::Infallible>> =
            vec![Ok(bad_embeddings)];

        let result = HnswIndex::build_batched_with_dim(batches.into_iter(), 3, 1024);
        assert!(
            result.is_err(),
            "build_batched_with_dim should reject dimension mismatch"
        );
        match result {
            Err(HnswError::DimensionMismatch { expected, actual }) => {
                assert_eq!(expected, 1024);
                assert_eq!(actual, 128);
            }
            Err(other) => panic!("Expected DimensionMismatch, got: {:?}", other),
            Ok(_) => panic!("Expected error, got Ok"),
        }
    }

    #[test]
    fn tc40_build_batched_with_dim_zero() {
        // dim=0 should return an error (not panic on chunks_exact(0))
        let batches: Vec<Result<Vec<(String, Embedding)>, std::convert::Infallible>> = vec![];
        let result = HnswIndex::build_batched_with_dim(batches.into_iter(), 0, 0);
        assert!(result.is_err(), "dim=0 should return error");
        let err = result.err().unwrap();
        assert!(
            err.to_string().contains("must be > 0"),
            "error message should mention dimension requirement"
        );
    }

    #[test]
    fn tc40_build_batched_with_dim_zero_nonempty_errors() {
        // dim=0 with actual embeddings should fail on dimension mismatch
        let embeddings: Vec<(String, Embedding)> = vec![("chunk1".to_string(), make_embedding(1))];
        let batches: Vec<Result<Vec<(String, Embedding)>, std::convert::Infallible>> =
            vec![Ok(embeddings)];
        let result = HnswIndex::build_batched_with_dim(batches.into_iter(), 1, 0);
        assert!(
            result.is_err(),
            "dim=0 with non-empty embeddings should error on dimension mismatch"
        );
    }

    /// A NaN vector interleaved between valid vectors must be skipped (it
    /// passes the zero-vector check because `NaN != 0.0` is true), the build
    /// must succeed, and the surviving ids must stay aligned with their
    /// insertion indices — the id_map desync class documented in the build
    /// loop.
    #[test]
    fn test_build_batched_skips_non_finite_embedding() {
        // One-hot vectors: cosine-orthogonal, so each query's exact match
        // is unambiguous even at k=1.
        fn one_hot(idx: usize) -> Embedding {
            let mut v = vec![0.0f32; crate::EMBEDDING_DIM];
            v[idx] = 1.0;
            Embedding::new(v)
        }

        // Non-zero NaN component: passes the zero-vector check, must be
        // caught by the non-finite check.
        let mut nan_vec = vec![0.0f32; crate::EMBEDDING_DIM];
        nan_vec[1] = f32::NAN;

        let batch = vec![
            ("chunk_a".to_string(), one_hot(0)),
            ("chunk_nan".to_string(), Embedding::new(nan_vec)),
            ("chunk_b".to_string(), one_hot(2)),
        ];
        let batches: Vec<Result<Vec<(String, Embedding)>, std::convert::Infallible>> =
            vec![Ok(batch)];

        let index = HnswIndex::build_batched_with_dim(batches.into_iter(), 3, crate::EMBEDDING_DIM)
            .expect("build must succeed with a non-finite vector in the stream");
        assert_eq!(index.len(), 2, "NaN vector must be skipped, not inserted");

        // Exact-embedding searches must return the correct ids — pins that
        // the skip kept id_map positions aligned with insertion indices.
        let hits = index.search(&one_hot(0), 1);
        assert_eq!(hits.first().map(|r| r.id.as_str()), Some("chunk_a"));
        let hits = index.search(&one_hot(2), 1);
        assert_eq!(hits.first().map(|r| r.id.as_str()), Some("chunk_b"));
    }

    /// A zero vector interleaved between valid vectors must be skipped via the
    /// zero-vector branch (distinct from the non-finite branch), the build must
    /// succeed, and the surviving ids must stay aligned with their insertion
    /// indices. Companion to `test_build_batched_skips_non_finite_embedding`,
    /// pinning the *zero-vector* arm of the id_map desync class specifically.
    #[test]
    fn test_build_batched_skips_interleaved_zero_vector() {
        // One-hot vectors: cosine-orthogonal, so each query's exact match is
        // unambiguous even at k=1.
        fn one_hot(idx: usize) -> Embedding {
            let mut v = vec![0.0f32; crate::EMBEDDING_DIM];
            v[idx] = 1.0;
            Embedding::new(v)
        }

        // All-zero vector: caught by the zero-vector check (every element 0.0).
        let zero_vec = vec![0.0f32; crate::EMBEDDING_DIM];

        let batch = vec![
            ("chunk_a".to_string(), one_hot(0)),
            ("chunk_zero".to_string(), Embedding::new(zero_vec)),
            ("chunk_b".to_string(), one_hot(2)),
        ];
        let batches: Vec<Result<Vec<(String, Embedding)>, std::convert::Infallible>> =
            vec![Ok(batch)];

        let index = HnswIndex::build_batched_with_dim(batches.into_iter(), 3, crate::EMBEDDING_DIM)
            .expect("build must succeed with a zero vector in the stream");
        assert_eq!(index.len(), 2, "zero vector must be skipped, not inserted");

        // Exact-embedding searches must return the correct ids — proves the
        // interleaved zero-vector skip kept the id_map aligned with insertion
        // indices (a desync would attribute chunk_a's vector to chunk_b's id).
        let hits = index.search(&one_hot(0), 1);
        assert_eq!(hits.first().map(|r| r.id.as_str()), Some("chunk_a"));
        let hits = index.search(&one_hot(2), 1);
        assert_eq!(hits.first().map(|r| r.id.as_str()), Some("chunk_b"));
    }

    /// HNSW search with `k` far larger than the index size must clamp to the
    /// index and return only the real vectors — no phantom or duplicate ids.
    /// Mirrors CAGRA's `test_search_k_greater_than_len_drops_phantoms`; the
    /// HNSW guard is the `ef_search.max(k*2).min(index_size)` clamp.
    #[test]
    fn test_search_k_greater_than_index_size() {
        let embeddings = vec![
            ("chunk_0".to_string(), make_embedding(0)),
            ("chunk_1".to_string(), make_embedding(1)),
            ("chunk_2".to_string(), make_embedding(2)),
        ];
        let index = HnswIndex::build_with_dim(embeddings, crate::EMBEDDING_DIM).unwrap();
        assert_eq!(index.len(), 3);

        let results = index.search(&make_embedding(1), 100);

        // Only three real vectors exist; the over-large k must not invent more.
        assert!(
            results.len() <= 3,
            "expected at most 3 results, got {}",
            results.len()
        );
        // Every id is a real chunk and the set is unique (no slot returned
        // twice from the clamped search).
        let mut seen = std::collections::HashSet::new();
        for r in &results {
            assert!(
                matches!(r.id.as_str(), "chunk_0" | "chunk_1" | "chunk_2"),
                "unexpected id in oversized-k search: {}",
                r.id
            );
            assert!(seen.insert(r.id.clone()), "duplicate id returned: {}", r.id);
        }
    }

    // ===== HnswIndex::search_filtered predicate tests =====

    #[test]
    fn tc17_search_filtered_accepts_only_matching_ids() {
        // Build a 20-item index where IDs are "chunk_0" .. "chunk_19".
        // Filter accepts only even-numbered IDs. Assert no odd IDs appear.
        let embeddings: Vec<(String, Embedding)> = (0..20)
            .map(|i| (format!("chunk_{}", i), make_embedding(i)))
            .collect();

        let index = HnswIndex::build_with_dim(embeddings, crate::EMBEDDING_DIM).unwrap();
        assert_eq!(index.len(), 20);

        let query = make_embedding(4); // bias toward chunk_4 (even)
        let results = index.search_filtered(&query, 10, &|id: &str| {
            // Accept only even-numbered chunk IDs
            id.strip_prefix("chunk_")
                .and_then(|n| n.parse::<u32>().ok())
                .is_some_and(|n| n % 2 == 0)
        });

        // Every returned result must have an even ID
        for r in &results {
            let num: u32 = r.id.strip_prefix("chunk_").unwrap().parse().unwrap();
            assert_eq!(
                num % 2,
                0,
                "Rejected ID '{}' appeared in search_filtered results",
                r.id
            );
        }
        // Should have returned some results (we have 10 even IDs, asked for 10)
        assert!(
            !results.is_empty(),
            "search_filtered should return results when matches exist"
        );
    }

    #[test]
    fn tc17_search_filtered_all_rejected_returns_empty() {
        // Build a small index, then filter that rejects everything.
        let embeddings: Vec<(String, Embedding)> = (0..10)
            .map(|i| (format!("chunk_{}", i), make_embedding(i)))
            .collect();

        let index = HnswIndex::build_with_dim(embeddings, crate::EMBEDDING_DIM).unwrap();

        let query = make_embedding(1);
        let results = index.search_filtered(&query, 5, &|_id: &str| false);

        assert!(
            results.is_empty(),
            "search_filtered with all-reject predicate should return empty, got {} results",
            results.len()
        );
    }

    /// `HnswIndex::search` with `k=0` must return an empty result without
    /// panicking. Pins the contract at the cqs boundary so a future
    /// `hnsw_rs` version bump can't silently change behavior — the internal
    /// `ef_search = self.ef_search.max(k * 2).min(index_size)` math must not
    /// trip on `k=0`.
    #[test]
    fn tc18_search_k_zero_returns_empty() {
        // Build a populated index so the k=0 path doesn't short-circuit on
        // `id_map.is_empty()` — we want to drive into search_neighbours.
        let embeddings: Vec<(String, Embedding)> = (0..10)
            .map(|i| (format!("chunk_{}", i), make_embedding(i)))
            .collect();
        let index = HnswIndex::build_with_dim(embeddings, crate::EMBEDDING_DIM).unwrap();

        let query = make_embedding(1);
        let results = index.search(&query, 0);
        assert!(
            results.is_empty(),
            "search with k=0 must return empty, got {} results",
            results.len()
        );
    }

    /// Empty-index direct test. The `id_map.is_empty()` short-circuit must
    /// apply to unfiltered `search()` too, not only the filtered path.
    #[test]
    fn tc_hap_search_empty_index_returns_empty() {
        let index = HnswIndex::build_with_dim(Vec::new(), crate::EMBEDDING_DIM).unwrap();
        let query = make_embedding(0);
        assert!(index.search(&query, 5).is_empty());
    }

    /// Dim-mismatch query must return empty without crashing the dense
    /// library. Pins the early-return in hnsw/search.rs.
    #[test]
    fn tc_hap_search_dim_mismatch_returns_empty() {
        let embeddings: Vec<(String, Embedding)> = (0..5)
            .map(|i| (format!("chunk_{}", i), make_embedding(i)))
            .collect();
        let index = HnswIndex::build_with_dim(embeddings, crate::EMBEDDING_DIM).unwrap();
        let bad = Embedding::new(vec![0.5; 32]); // wrong dim
        assert!(index.search(&bad, 5).is_empty());
    }

    /// NaN/Inf query must short-circuit before the `assert dist >= ε`
    /// panic in the dense library.
    #[test]
    fn tc_hap_search_nonfinite_query_returns_empty() {
        let embeddings: Vec<(String, Embedding)> = (0..5)
            .map(|i| (format!("chunk_{}", i), make_embedding(i)))
            .collect();
        let index = HnswIndex::build_with_dim(embeddings, crate::EMBEDDING_DIM).unwrap();
        let nan_query = Embedding::new(vec![f32::NAN; crate::EMBEDDING_DIM]);
        assert!(index.search(&nan_query, 5).is_empty());
    }

    /// The DotProduct arm must build AND self-query without panicking on
    /// f32-L2-normalized vectors. A vector L2-normalized in f32 has a self-dot
    /// fractionally above 1.0 (≈ 1.0000005); the stock dot distance asserts
    /// `a·b <= 1` and would unwind a parallel-insert worker on the build path
    /// (killing the whole build) and panic on the self-match search path. The
    /// clamped dot distance maps that to distance 0 instead. Sweeping a wide
    /// seed band exercises the many seeds that produce a >1.0 self-dot.
    #[test]
    fn dot_build_and_self_query_does_not_panic_over_seed_band() {
        use crate::index::DistanceMetric;
        // Build one DotProduct index over a band of make_test_embedding seeds,
        // then self-query EACH indexed vector. The previous failure unwound
        // during parallel_insert_data (build) or during search_neighbours
        // (the self-match dot > 1.0), so reaching the end without a panic IS
        // the assertion.
        let embeddings: Vec<(String, Embedding)> = (0..64u32)
            .map(|i| (format!("chunk_{i}"), make_embedding(i)))
            .collect();
        let index = HnswIndex::build_with_dim_and_metric(
            embeddings,
            crate::EMBEDDING_DIM,
            DistanceMetric::DotProduct,
        )
        .expect("DotProduct build must not panic on f32-normalized vectors");
        assert_eq!(index.len(), 64);

        for i in 0..64u32 {
            // The query is an indexed vector; its self-dot is the >1.0 input
            // that previously tripped the assert. Search must return finite,
            // in-range scores rather than panic.
            let q = make_embedding(i);
            let hits = index.search(&q, 5);
            for h in &hits {
                assert!(
                    h.score.is_finite(),
                    "non-finite dot score for seed {i}: {}",
                    h.score
                );
                // score = 1 - dist = min(a·b, 1) — never exceeds 1.0.
                assert!(
                    h.score <= 1.0 + 1e-5,
                    "clamped dot score must not exceed 1.0 for seed {i}: {}",
                    h.score
                );
            }
        }
    }
}
