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

fn build_index() -> AnnIndex {
    let mut index = AnnIndex::new(AnnConfig::for_dimensions(4)).unwrap();
    let records: Vec<(String, Vec<f32>)> = (0..500)
        .map(|i| {
            let angle = i as f32 * 0.041;
            (
                format!("record-{i:05}"),
                vec![
                    angle.sin(),
                    angle.cos(),
                    (angle * 3.0).sin(),
                    (angle * 3.0).cos(),
                ],
            )
        })
        .collect();
    index.bulk_load(records, 4).unwrap();
    index.delete("record-00042");
    index
}

fn probe(index: &AnnIndex, seed: f32) -> Vec<crate::AnnCandidate> {
    let angle = seed * 0.29;
    index
        .query(AnnQuery {
            vector: &[
                angle.sin(),
                angle.cos(),
                (angle * 3.0).sin(),
                (angle * 3.0).cos(),
            ],
            top_k: 10,
            ef_search: Some(64),
            filter: None,
        })
        .unwrap()
}

#[test]
fn snapshot_roundtrip_preserves_the_graph() {
    let index = build_index();
    let mut bytes = Vec::new();
    index.write_to(&mut bytes).unwrap();

    let restored = AnnIndex::read_from(bytes.as_slice()).unwrap();
    assert_eq!(restored.len(), index.len());
    for i in 0..25 {
        assert_eq!(
            probe(&index, i as f32),
            probe(&restored, i as f32),
            "snapshot roundtrip must preserve query results exactly"
        );
    }
}

#[test]
fn snapshot_roundtrip_preserves_tombstones() {
    let index = build_index();
    let mut bytes = Vec::new();
    index.write_to(&mut bytes).unwrap();
    let mut restored = AnnIndex::read_from(bytes.as_slice()).unwrap();

    assert!(!restored.delete("record-00042"));
    assert!(restored.delete("record-00043"));
}

#[test]
fn corrupted_snapshots_are_rejected() {
    let index = build_index();
    let mut bytes = Vec::new();
    index.write_to(&mut bytes).unwrap();

    for position in [0, 8, bytes.len() / 2, bytes.len() - 5] {
        let mut corrupted = bytes.clone();
        corrupted[position] ^= 0x40;
        assert!(
            AnnIndex::read_from(corrupted.as_slice()).is_err(),
            "flipped byte at {position} must be detected"
        );
    }

    let truncated = &bytes[..bytes.len() - 12];
    assert!(AnnIndex::read_from(truncated).is_err());
}

#[test]
fn snapshot_saves_and_loads_from_disk() {
    let index = build_index();
    let path = std::env::temp_dir().join(format!("feox-ann-test-{}.fxa", std::process::id()));
    index.save_to(&path).unwrap();
    let restored = AnnIndex::load_from(&path).unwrap();
    std::fs::remove_file(&path).ok();

    assert_eq!(restored.len(), index.len());
    assert_eq!(probe(&index, 1.0), probe(&restored, 1.0));
}
