//! Tests for `BruteForceIndex`.

#![allow(clippy::unwrap_used)]

use super::{
    super::{DistanceMetric, QuantizationLevel},
    *,
};

/// Helper: create a 3-dimensional config for compact tests.
fn test_config() -> VectorIndexConfig {
    VectorIndexConfig {
        dimensions: 3,
        metric: DistanceMetric::L2,
        quantization: QuantizationLevel::F32,
        connectivity: 16,
        ef_construction: 200,
        ef_search: 200,
    }
}

// -----------------------------------------------------------------------
// Unit tests
// -----------------------------------------------------------------------

#[test]
fn brute_force_new_is_empty() {
    let index = BruteForceIndex::new(test_config());
    assert!(index.is_empty());
    assert_eq!(index.len(), 0);
    assert_eq!(index.len_live(), 0);
}

#[test]
fn brute_force_insert_and_len() {
    let index = BruteForceIndex::new(test_config());
    index.insert(1, &[1.0, 2.0, 3.0]).unwrap();
    index.insert(2, &[4.0, 5.0, 6.0]).unwrap();
    assert_eq!(index.len(), 2);
    assert_eq!(index.len_live(), 2);
    assert!(!index.is_empty());
}

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

#[test]
fn brute_force_search_dimension_mismatch() {
    let index = BruteForceIndex::new(test_config());
    index.insert(1, &[1.0, 2.0, 3.0]).unwrap();
    let result = index.search(&[1.0, 2.0], 1, 200); // 2D query on 3D index
    assert!(result.is_err());
    assert!(matches!(
        result.unwrap_err(),
        VectorError::DimensionMismatch {
            expected: 3,
            got: 2
        }
    ));
}

#[test]
fn brute_force_self_search_distance_zero() {
    let index = BruteForceIndex::new(test_config());
    let vec = [1.0, 2.0, 3.0];
    index.insert(1, &vec).unwrap();

    let results = index.search(&vec, 1, 200).unwrap();
    assert_eq!(results.len(), 1);
    assert_eq!(results[0].id, 1);
    assert!((results[0].distance - 0.0).abs() < f32::EPSILON);
}

#[test]
fn brute_force_search_empty_index() {
    let index = BruteForceIndex::new(test_config());
    let results = index.search(&[1.0, 2.0, 3.0], 5, 200).unwrap();
    assert!(results.is_empty());
}

#[test]
fn brute_force_search_k_larger_than_index() {
    let index = BruteForceIndex::new(test_config());
    index.insert(1, &[1.0, 0.0, 0.0]).unwrap();
    index.insert(2, &[0.0, 1.0, 0.0]).unwrap();

    let results = index.search(&[1.0, 0.0, 0.0], 10, 200).unwrap();
    assert_eq!(results.len(), 2); // only 2 vectors exist
}

#[test]
fn brute_force_small_k_returns_correct_top_k() {
    // Exercises the bounded-selection path (k << n): select_nth_unstable_by must
    // return exactly the k nearest, in ascending-distance order, identical to a
    // full sort. Build 8 vectors at increasing distance along the x-axis from the
    // query [0,0,0]; the k nearest are the k smallest x magnitudes.
    let index = BruteForceIndex::new(test_config());
    for i in 1..=8u64 {
        #[allow(clippy::cast_precision_loss)]
        let x = i as f32; // distances are x^2: 1,4,9,...,64 — all distinct
        index.insert(i, &[x, 0.0, 0.0]).unwrap();
    }

    let results = index.search(&[0.0, 0.0, 0.0], 3, 200).unwrap();
    assert_eq!(results.len(), 3, "k=3 over 8 vectors");
    // Nearest three are ids 1,2,3 (x = 1,2,3) in that order.
    assert_eq!(results[0].id, 1);
    assert_eq!(results[1].id, 2);
    assert_eq!(results[2].id, 3);
    // Ascending distance ordering preserved.
    for pair in results.windows(2) {
        assert!(pair[0].distance <= pair[1].distance);
    }
}

#[test]
fn brute_force_small_k_with_tied_distances() {
    // Ties must not corrupt the partition: with several equal nearest distances,
    // top-k still returns k results all at the minimum distance.
    let index = BruteForceIndex::new(test_config());
    // Four vectors all at distance 1 from query, two further out.
    index.insert(1, &[1.0, 0.0, 0.0]).unwrap();
    index.insert(2, &[-1.0, 0.0, 0.0]).unwrap();
    index.insert(3, &[0.0, 1.0, 0.0]).unwrap();
    index.insert(4, &[0.0, -1.0, 0.0]).unwrap();
    index.insert(5, &[3.0, 0.0, 0.0]).unwrap(); // distance 9
    index.insert(6, &[0.0, 4.0, 0.0]).unwrap(); // distance 16

    let results = index.search(&[0.0, 0.0, 0.0], 2, 200).unwrap();
    assert_eq!(results.len(), 2);
    // Both returned must be among the four unit-distance vectors.
    for r in &results {
        assert!(
            (r.distance - 1.0).abs() < f32::EPSILON,
            "expected tied min dist"
        );
        assert!((1..=4).contains(&r.id), "id {} not in tied set", r.id);
    }
}

#[test]
fn brute_force_zero_k_returns_empty() {
    let index = BruteForceIndex::new(test_config());
    index.insert(1, &[1.0, 0.0, 0.0]).unwrap();
    let results = index.search(&[1.0, 0.0, 0.0], 0, 200).unwrap();
    assert!(results.is_empty());
}

#[test]
fn brute_force_orthogonal_vectors_distance() {
    let index = BruteForceIndex::new(test_config());
    // Two unit vectors along different axes: L2^2 = 2.0
    index.insert(1, &[1.0, 0.0, 0.0]).unwrap();
    index.insert(2, &[0.0, 1.0, 0.0]).unwrap();

    let results = index.search(&[1.0, 0.0, 0.0], 2, 200).unwrap();
    assert_eq!(results[0].id, 1); // self is closest
    assert!((results[0].distance - 0.0).abs() < f32::EPSILON);
    assert_eq!(results[1].id, 2);
    assert!((results[1].distance - 2.0).abs() < f32::EPSILON);
}

#[test]
fn brute_force_identical_vectors_distance() {
    let index = BruteForceIndex::new(test_config());
    index.insert(1, &[3.0, 4.0, 5.0]).unwrap();
    index.insert(2, &[3.0, 4.0, 5.0]).unwrap();

    let results = index.search(&[3.0, 4.0, 5.0], 2, 200).unwrap();
    assert_eq!(results.len(), 2);
    assert!((results[0].distance - 0.0).abs() < f32::EPSILON);
    assert!((results[1].distance - 0.0).abs() < f32::EPSILON);
}

#[test]
fn brute_force_delete_and_search() {
    let index = BruteForceIndex::new(test_config());
    index.insert(1, &[1.0, 0.0, 0.0]).unwrap();
    index.insert(2, &[0.0, 1.0, 0.0]).unwrap();
    index.insert(3, &[0.0, 0.0, 1.0]).unwrap();

    index.delete(2).unwrap();

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

#[test]
fn brute_force_delete_not_found() {
    let index = BruteForceIndex::new(test_config());
    let result = index.delete(999);
    assert!(result.is_err());
    assert!(matches!(
        result.unwrap_err(),
        VectorError::NotFound { id: 999 }
    ));
}

#[test]
fn brute_force_insert_replaces_existing() {
    let index = BruteForceIndex::new(test_config());
    index.insert(1, &[1.0, 0.0, 0.0]).unwrap();
    // Replace vector 1 with a different embedding
    index.insert(1, &[0.0, 0.0, 1.0]).unwrap();

    assert_eq!(index.len(), 1); // still just one vector
    let results = index.search(&[0.0, 0.0, 1.0], 1, 200).unwrap();
    assert_eq!(results[0].id, 1);
    assert!((results[0].distance - 0.0).abs() < f32::EPSILON);
}

#[test]
fn brute_force_filtered_search_excludes_non_matching() {
    let index = BruteForceIndex::new(test_config());
    index.insert(1, &[1.0, 0.0, 0.0]).unwrap();
    index.insert(2, &[0.0, 1.0, 0.0]).unwrap();
    index.insert(3, &[0.0, 0.0, 1.0]).unwrap();

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

    assert_eq!(results.len(), 1);
    assert_eq!(results[0].id, 2);
}

#[test]
fn brute_force_filtered_search_empty_result() {
    let index = BruteForceIndex::new(test_config());
    index.insert(1, &[1.0, 0.0, 0.0]).unwrap();
    index.insert(3, &[0.0, 0.0, 1.0]).unwrap();

    // Filter rejects all (only even IDs, but we only have odd)
    let results = index
        .filtered_search(&[1.0, 0.0, 0.0], 3, 200, &|id| id % 2 == 0)
        .unwrap();

    assert!(results.is_empty());
}

#[test]
fn brute_force_save_load_roundtrip() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("test_index.bfvi");

    let config = test_config();
    let index = BruteForceIndex::new(config.clone());
    index.insert(1, &[1.0, 2.0, 3.0]).unwrap();
    index.insert(2, &[4.0, 5.0, 6.0]).unwrap();
    index.insert(3, &[7.0, 8.0, 9.0]).unwrap();

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

    let loaded = BruteForceIndex::load(&path, &config).unwrap();
    assert_eq!(loaded.len(), 3);

    // Verify search produces the same results
    let query = [1.0, 2.0, 3.0];
    let original_results = index.search(&query, 3, 200).unwrap();
    let loaded_results = loaded.search(&query, 3, 200).unwrap();

    assert_eq!(original_results.len(), loaded_results.len());
    for (orig, load) in original_results.iter().zip(loaded_results.iter()) {
        assert_eq!(orig.id, load.id);
        assert!((orig.distance - load.distance).abs() < f32::EPSILON);
    }
}

#[test]
fn brute_force_reserve_is_noop() {
    let index = BruteForceIndex::new(test_config());
    // Must not error
    index.reserve(1_000_000).unwrap();
    assert!(index.is_empty()); // no side effects
}

#[test]
fn l2_distance_sq_correctness() {
    // [1, 2, 3] vs [4, 5, 6] => (3^2 + 3^2 + 3^2) = 27
    let dist = l2_distance_sq(&[1.0, 2.0, 3.0], &[4.0, 5.0, 6.0]);
    assert!((dist - 27.0).abs() < f32::EPSILON);

    // Identical vectors => 0
    let dist = l2_distance_sq(&[1.0, 2.0, 3.0], &[1.0, 2.0, 3.0]);
    assert!((dist - 0.0).abs() < f32::EPSILON);

    // Single dimension
    let dist = l2_distance_sq(&[3.0], &[7.0]);
    assert!((dist - 16.0).abs() < f32::EPSILON);
}

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

#[test]
fn vector_index_config_defaults() {
    let config = VectorIndexConfig::default();
    assert_eq!(config.dimensions, 1536);
    assert_eq!(config.metric, DistanceMetric::L2);
    assert_eq!(config.quantization, QuantizationLevel::F16);
    assert_eq!(config.connectivity, 16);
    assert_eq!(config.ef_construction, 400);
    assert_eq!(config.ef_search, 200);
}

// -----------------------------------------------------------------------
// Corruption / hostile-input deserialization
// -----------------------------------------------------------------------

/// Build a minimal valid 17-byte brute-force index header for `dims`
/// dimensions claiming `count` stored vectors. The returned buffer contains
/// *only* the header -- callers append (or omit) body bytes to construct
/// truncated or torn inputs.
fn make_header(dims: u32, count: u64) -> Vec<u8> {
    let mut buf = Vec::with_capacity(17);
    buf.extend_from_slice(MAGIC); // [0..4]
    buf.push(FORMAT_VERSION); // [4]
    buf.extend_from_slice(&dims.to_le_bytes()); // [5..9]
    buf.extend_from_slice(&count.to_le_bytes()); // [9..17]
    buf
}

#[test]
fn deserialize_rejects_enormous_count_without_panicking() {
    // A header claiming u64::MAX vectors with no body. Pre-fix this wrapped
    // `count * bytes_per_vector`, slipped past the truncation guard, then
    // aborted in `HashMap::with_capacity` / panicked in the read loop.
    let config = test_config(); // dims = 3
    let data = make_header(config.dimensions as u32, u64::MAX);

    let result = BruteForceIndex::deserialize(&data, &config);
    assert!(
        matches!(result, Err(VectorError::CorruptedIndex(_))),
        "enormous count must be a recoverable CorruptedIndex, got {result:?}"
    );
}

#[test]
fn deserialize_rejects_count_exceeding_body() {
    // Header claims 1000 vectors but supplies a body big enough for only one.
    let config = test_config(); // dims = 3 => bytes_per_vector = 8 + 12 = 20
    let mut data = make_header(config.dimensions as u32, 1000);
    data.extend_from_slice(&[0u8; 20]); // exactly one vector's worth of body

    let result = BruteForceIndex::deserialize(&data, &config);
    match result {
        Err(VectorError::CorruptedIndex(msg)) => {
            assert!(
                msg.contains("implausible vector count"),
                "expected count-cap message, got: {msg}"
            );
        }
        other => panic!("expected CorruptedIndex, got {other:?}"),
    }
}

#[test]
fn deserialize_rejects_torn_body() {
    // Header honestly claims 2 vectors (40 bytes of body) but the body is torn
    // mid-vector (only 25 bytes present). This must be rejected with a typed
    // CorruptedIndex error, never read OOB. The count-cap guard catches it first:
    // a 25-byte body holds at most 1 whole 20-byte vector, so a claimed count of 2
    // is already implausible — which is exactly the bound that prevents the OOB
    // read the torn tail would otherwise cause.
    let config = test_config(); // dims = 3 => bytes_per_vector = 20
    let mut data = make_header(config.dimensions as u32, 2);
    data.extend_from_slice(&[0u8; 25]); // 25 < 40: torn second vector

    let result = BruteForceIndex::deserialize(&data, &config);
    match result {
        Err(VectorError::CorruptedIndex(msg)) => {
            assert!(
                msg.contains("implausible vector count") || msg.contains("truncated"),
                "expected a size/count corruption message, got: {msg}"
            );
        }
        other => panic!("expected CorruptedIndex on torn body, got {other:?}"),
    }
}

#[test]
fn deserialize_rejects_overflowing_dimensions() {
    // A dims value whose `* 4` overflows `usize` must be rejected, not panic.
    // We point `config.dimensions` at the same huge value so the dimension
    // check passes and we reach the per-vector size arithmetic.
    let huge_dims = usize::MAX;
    let config = VectorIndexConfig {
        dimensions: huge_dims,
        ..test_config()
    };
    // Header stores dims as a u32, so truncation means the on-file dims won't
    // equal `huge_dims`; that mismatch is itself a CorruptedIndex, which is
    // still the correct (non-panicking) outcome. Assert no panic + an error.
    let data = make_header(huge_dims as u32, 1);
    let result = BruteForceIndex::deserialize(&data, &config);
    assert!(
        matches!(result, Err(VectorError::CorruptedIndex(_))),
        "overflowing dims must yield CorruptedIndex, got {result:?}"
    );
}

#[test]
fn deserialize_accepts_well_formed_buffer() {
    // Round-trip sanity: a buffer produced by `save` with a real body still
    // deserializes after the hardening. Guards the fix against over-rejection.
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("ok.bfvi");
    let config = test_config();
    let index = BruteForceIndex::new(config.clone());
    index.insert(7, &[1.0, 2.0, 3.0]).unwrap();
    index.insert(8, &[4.0, 5.0, 6.0]).unwrap();
    index.save(&path).unwrap();

    let data = std::fs::read(&path).unwrap();
    let loaded = BruteForceIndex::deserialize(&data, &config).unwrap();
    assert_eq!(loaded.len(), 2);
}

// -----------------------------------------------------------------------
// Property tests
// -----------------------------------------------------------------------

#[allow(clippy::cast_precision_loss)]
mod proptests {
    use proptest::prelude::*;

    use super::*;

    /// Generate a deterministic vector from an ID and dimension index.
    /// This avoids needing `rand` -- the values are fully determined by
    /// the ID and dimension.
    fn deterministic_vector(id: u64, dims: usize) -> Vec<f32> {
        (0..dims)
            .map(|i| ((id as usize * 3 + i * 7) % 100) as f32 / 100.0 - 0.5)
            .collect()
    }

    /// Normalize a vector to unit length. Returns the zero vector unchanged
    /// (property tests generate non-zero vectors via the arithmetic above).
    fn normalize(v: &mut [f32]) {
        let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        if norm > f32::EPSILON {
            for x in v.iter_mut() {
                *x /= norm;
            }
        }
    }

    proptest! {
        #[test]
        fn insert_search_roundtrip(count in 1_usize..50) {
            let dims = 8;
            let config = VectorIndexConfig {
                dimensions: dims,
                metric: DistanceMetric::L2,
                quantization: QuantizationLevel::F32,
                connectivity: 16,
                ef_construction: 200,
                ef_search: 200,
            };
            let index = BruteForceIndex::new(config);

            // Insert deterministic vectors
            for id in 0..count as u64 {
                let mut vec = deterministic_vector(id, dims);
                normalize(&mut vec);
                index.insert(id, &vec).unwrap();
            }

            // Each vector should be its own nearest neighbor (distance ~ 0)
            for id in 0..count as u64 {
                let mut query = deterministic_vector(id, dims);
                normalize(&mut query);

                let results = index.search(&query, 1, 200).unwrap();
                prop_assert!(!results.is_empty());

                // The vector should find itself. Access the stored vector
                // to verify the distance is indeed ~0.
                let stored: Vec<f32> = index.vectors.read().unwrap()[&id].clone();
                let dist = l2_distance_sq(&query, &stored);
                prop_assert!(dist < 1e-6, "self-distance should be ~0, got {}", dist);
            }
        }

        #[test]
        fn delete_excludes_from_results(count in 2_usize..50) {
            let dims = 8;
            let config = VectorIndexConfig {
                dimensions: dims,
                metric: DistanceMetric::L2,
                quantization: QuantizationLevel::F32,
                connectivity: 16,
                ef_construction: 200,
                ef_search: 200,
            };
            let index = BruteForceIndex::new(config);

            for id in 0..count as u64 {
                let vec = deterministic_vector(id, dims);
                index.insert(id, &vec).unwrap();
            }

            // Delete the first vector
            index.delete(0).unwrap();

            // Search for all remaining vectors -- ID 0 must never appear
            let query = deterministic_vector(0, dims);
            let results = index.search(&query, count, 200).unwrap();

            for result in &results {
                prop_assert_ne!(result.id, 0, "deleted vector 0 found in results");
            }
            prop_assert_eq!(results.len(), count - 1);
        }

        #[test]
        fn filtered_search_honors_predicate(count in 1_usize..50) {
            let dims = 8;
            let config = VectorIndexConfig {
                dimensions: dims,
                metric: DistanceMetric::L2,
                quantization: QuantizationLevel::F32,
                connectivity: 16,
                ef_construction: 200,
                ef_search: 200,
            };
            let index = BruteForceIndex::new(config);

            for id in 0..count as u64 {
                let vec = deterministic_vector(id, dims);
                index.insert(id, &vec).unwrap();
            }

            // Filter: only even IDs
            let query = deterministic_vector(0, dims);
            let results = index
                .filtered_search(&query, count, 200, &|id| id % 2 == 0)
                .unwrap();

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

        #[test]
        fn results_sorted_by_distance(count in 1_usize..50) {
            let dims = 8;
            let config = VectorIndexConfig {
                dimensions: dims,
                metric: DistanceMetric::L2,
                quantization: QuantizationLevel::F32,
                connectivity: 16,
                ef_construction: 200,
                ef_search: 200,
            };
            let index = BruteForceIndex::new(config);

            for id in 0..count as u64 {
                let vec = deterministic_vector(id, dims);
                index.insert(id, &vec).unwrap();
            }

            let query = deterministic_vector(0, dims);
            let results = index.search(&query, count, 200).unwrap();

            // Verify ascending distance ordering
            for pair in results.windows(2) {
                prop_assert!(
                    pair[0].distance <= pair[1].distance,
                    "results not sorted: {} > {}",
                    pair[0].distance,
                    pair[1].distance
                );
            }
        }
    }
}
