use crate::{AnnConfig, AnnIndex, AnnQuery};

#[test]
fn queries_nearest_vectors() {
    let mut index = AnnIndex::new(AnnConfig::for_dimensions(2)).unwrap();
    index.upsert("north".to_string(), &[1.0, 0.0]).unwrap();
    index.upsert("east".to_string(), &[0.0, 1.0]).unwrap();
    index.upsert("south".to_string(), &[-1.0, 0.0]).unwrap();

    let matches = index
        .query(AnnQuery {
            vector: &[0.9, 0.1],
            top_k: 2,
            ef_search: Some(8),
            filter: None,
        })
        .unwrap();

    assert_eq!(matches[0].id, "north");
    assert!(matches[0].score > matches[1].score);
}

#[test]
fn delete_hides_records_from_results() {
    let mut index = AnnIndex::new(AnnConfig::for_dimensions(2)).unwrap();
    index.upsert("north".to_string(), &[1.0, 0.0]).unwrap();
    index.upsert("east".to_string(), &[0.0, 1.0]).unwrap();

    assert!(index.delete("north"));
    let matches = index
        .query(AnnQuery {
            vector: &[1.0, 0.0],
            top_k: 2,
            ef_search: Some(8),
            filter: None,
        })
        .unwrap();

    assert_eq!(matches.len(), 1);
    assert_eq!(matches[0].id, "east");
}

#[test]
fn upsert_replaces_visible_record() {
    let mut index = AnnIndex::new(AnnConfig::for_dimensions(2)).unwrap();
    index.upsert("item".to_string(), &[1.0, 0.0]).unwrap();
    index.upsert("item".to_string(), &[0.0, 1.0]).unwrap();

    let matches = index
        .query(AnnQuery {
            vector: &[0.0, 1.0],
            top_k: 1,
            ef_search: Some(8),
            filter: None,
        })
        .unwrap();

    assert_eq!(index.len(), 1);
    assert_eq!(matches[0].id, "item");
    assert!(matches[0].score > 0.99);
}

#[test]
fn query_applies_filter_to_results() {
    let mut index = AnnIndex::new(AnnConfig::for_dimensions(2)).unwrap();
    index.upsert("north".to_string(), &[1.0, 0.0]).unwrap();
    index.upsert("north-east".to_string(), &[0.8, 0.2]).unwrap();
    index.upsert("east".to_string(), &[0.0, 1.0]).unwrap();

    let filter = |id: &str| id == "north-east" || id == "east";
    let matches = index
        .query(AnnQuery {
            vector: &[1.0, 0.0],
            top_k: 2,
            ef_search: Some(8),
            filter: Some(&filter),
        })
        .unwrap();

    assert_eq!(matches[0].id, "north-east");
    assert_eq!(matches[1].id, "east");
}

#[test]
fn rebuilds_are_deterministic() {
    let build = || {
        let mut index = AnnIndex::new(AnnConfig::for_dimensions(4)).unwrap();
        for i in 0..200 {
            let angle = i as f32 * 0.031;
            index
                .upsert(
                    format!("record-{i:04}"),
                    &[
                        angle.sin(),
                        angle.cos(),
                        (angle * 2.0).sin(),
                        (angle * 2.0).cos(),
                    ],
                )
                .unwrap();
        }
        index
    };

    let first = build();
    let second = build();

    for i in 0..20 {
        let angle = i as f32 * 0.17;
        let query = AnnQuery {
            vector: &[
                angle.sin(),
                angle.cos(),
                (angle * 2.0).sin(),
                (angle * 2.0).cos(),
            ],
            top_k: 10,
            ef_search: Some(32),
            filter: None,
        };
        let left = first.query(query).unwrap();
        let right = second.query(query).unwrap();
        assert_eq!(
            left, right,
            "identical builds must return identical results"
        );
    }
}
