#![allow(clippy::unwrap_used)]

use criterion::{Criterion, black_box, criterion_group, criterion_main};
use tidaldb::query::{HybridFusion, RetrievalMode, route_results};
use tidaldb::schema::EntityId;

fn make_bm25(n: u64) -> Vec<(EntityId, f32)> {
    (0..n).map(|i| (EntityId::new(i), (n - i) as f32)).collect()
}

fn make_ann(n: u64) -> Vec<(EntityId, f32)> {
    (500..500 + n)
        .enumerate()
        .map(|(i, id)| (EntityId::new(id), i as f32 * 0.001))
        .collect()
}

fn bench_rrf_fuse_1k(c: &mut Criterion) {
    let bm25 = make_bm25(1_000);
    let ann = make_ann(1_000);
    let fusion = HybridFusion::new();

    c.bench_function("rrf_fuse_1k_per_list", |b| {
        b.iter(|| {
            let results = fusion.fuse(black_box(&bm25), black_box(&ann));
            black_box(results)
        });
    });
}

fn bench_route_hybrid(c: &mut Criterion) {
    let bm25 = make_bm25(1_000);
    let ann = make_ann(1_000);
    let fusion = HybridFusion::new();

    c.bench_function("route_hybrid_1k", |b| {
        b.iter(|| {
            let r = route_results(
                black_box(RetrievalMode::Hybrid),
                black_box(&bm25),
                black_box(&ann),
                black_box(&fusion),
            );
            black_box(r)
        });
    });
}

fn bench_route_text_only(c: &mut Criterion) {
    let bm25 = make_bm25(1_000);
    let fusion = HybridFusion::new();

    c.bench_function("route_text_only_1k", |b| {
        b.iter(|| {
            let r = route_results(
                black_box(RetrievalMode::TextOnly),
                black_box(&bm25),
                &[],
                black_box(&fusion),
            );
            black_box(r)
        });
    });
}

criterion_group!(
    fusion_benches,
    bench_rrf_fuse_1k,
    bench_route_hybrid,
    bench_route_text_only
);
criterion_main!(fusion_benches);