use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use feox_ann::{AnnConfig, AnnIndex, AnnQuery};

const DIMENSIONS: usize = 128;
const RECORDS: usize = 10_000;

struct XorShift(u64);

impl XorShift {
    fn next_f32(&mut self) -> f32 {
        self.0 ^= self.0 << 13;
        self.0 ^= self.0 >> 7;
        self.0 ^= self.0 << 17;
        ((self.0 >> 40) as f32 / (1 << 24) as f32) * 2.0 - 1.0
    }

    fn vector(&mut self, dimensions: usize) -> Vec<f32> {
        (0..dimensions).map(|_| self.next_f32()).collect()
    }
}

fn sample_records(records: usize) -> Vec<(String, Vec<f32>)> {
    let mut rng = XorShift(0x5eed_f0f5 ^ records as u64);
    (0..records)
        .map(|i| (format!("record-{i:06}"), rng.vector(DIMENSIONS)))
        .collect()
}

fn build_index(records: usize) -> AnnIndex {
    build_index_threaded(records, available_threads())
}

fn build_index_threaded(records: usize, threads: usize) -> AnnIndex {
    let mut index = AnnIndex::new(AnnConfig::for_dimensions(DIMENSIONS)).unwrap();
    index.bulk_load(sample_records(records), threads).unwrap();
    index
}

fn available_threads() -> usize {
    std::thread::available_parallelism()
        .map(|threads| threads.get())
        .unwrap_or(1)
}

fn bench_build(c: &mut Criterion) {
    let mut group = c.benchmark_group("build");
    group.sample_size(10);
    group.throughput(Throughput::Elements(RECORDS as u64));
    group.bench_function(BenchmarkId::new("bulk_load_1_thread", RECORDS), |b| {
        b.iter(|| build_index_threaded(RECORDS, 1));
    });
    group.bench_function(
        BenchmarkId::new(
            format!("bulk_load_{}_threads", available_threads()),
            RECORDS,
        ),
        |b| {
            b.iter(|| build_index_threaded(RECORDS, available_threads()));
        },
    );
    group.finish();
}

fn bench_query(c: &mut Criterion) {
    let index = build_index(RECORDS);
    let mut rng = XorShift(0xfeed_beef);
    let queries: Vec<Vec<f32>> = (0..256).map(|_| rng.vector(DIMENSIONS)).collect();

    let mut group = c.benchmark_group("query");
    group.throughput(Throughput::Elements(1));
    for ef_search in [32_usize, 64, 128, 256] {
        group.bench_function(BenchmarkId::new("top10_ef", ef_search), |b| {
            let mut cursor = 0;
            b.iter(|| {
                let query = &queries[cursor % queries.len()];
                cursor += 1;
                index
                    .query(AnnQuery {
                        vector: query,
                        top_k: 10,
                        ef_search: Some(ef_search),
                        filter: None,
                    })
                    .unwrap()
            });
        });
    }

    let filter = |id: &str| id.as_bytes().last().is_some_and(|byte| byte % 2 == 0);
    group.bench_function("top10_ef64_filtered", |b| {
        let mut cursor = 0;
        b.iter(|| {
            let query = &queries[cursor % queries.len()];
            cursor += 1;
            index
                .query(AnnQuery {
                    vector: query,
                    top_k: 10,
                    ef_search: Some(64),
                    filter: Some(&filter),
                })
                .unwrap()
        });
    });
    group.finish();
}

criterion_group!(benches, bench_build, bench_query);
criterion_main!(benches);
