use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
use tici_shard::{compactor::frag_utils::merge_row_entries_for_bench, utils::row_entry::RowEntry};

#[derive(Clone, Copy)]
struct BenchCase {
    name: &'static str,
    frag_count: usize,
    keys_per_frag: usize,
    duplicate_every: usize,
}

fn build_entries(case: BenchCase) -> Vec<Vec<RowEntry>> {
    let mut entries_list = Vec::with_capacity(case.frag_count);

    for frag_idx in 0..case.frag_count {
        let mut entries = Vec::with_capacity(case.keys_per_frag);
        for key_idx in 0..case.keys_per_frag {
            let base_key_id = key_idx * (case.frag_count + 1);
            let key_id = if case.duplicate_every > 0 && key_idx % case.duplicate_every == 0 {
                // Shared by all frags to simulate cross-frag overlap.
                base_key_id
            } else {
                // Mostly unique key space per frag.
                base_key_id + frag_idx + 1
            };

            let key = format!("k{:020}", key_id).into_bytes();
            let del_mark = (key_idx + frag_idx) % 11 == 0;
            entries.push(RowEntry::new(key, del_mark, 0));
        }
        entries_list.push(entries);
    }

    entries_list
}

fn bench_dedup_merge(c: &mut Criterion) {
    let cases = [
        BenchCase {
            name: "high_frag_low_dup",
            frag_count: 2048,
            keys_per_frag: 512,
            duplicate_every: 256,
        },
        BenchCase {
            name: "high_frag_sparse_dup",
            frag_count: 4096,
            keys_per_frag: 256,
            duplicate_every: 128,
        },
        BenchCase {
            name: "mid_frag_medium_dup",
            frag_count: 512,
            keys_per_frag: 1024,
            duplicate_every: 8,
        },
    ];

    let mut group = c.benchmark_group("dedup_merge");
    for case in cases {
        let entries_list = build_entries(case);
        let input_key_count: usize = entries_list.iter().map(Vec::len).sum();
        let (retained_key_count, obsolete_by_dedup_count) =
            merge_row_entries_for_bench(&entries_list).unwrap();
        println!(
            "case={}, frag_count={}, keys_per_frag={}, input_key_count={}, retained_key_count={}, obsolete_by_dedup_count={}",
            case.name,
            case.frag_count,
            case.keys_per_frag,
            input_key_count,
            retained_key_count,
            obsolete_by_dedup_count,
        );

        let bench_id = BenchmarkId::new(
            case.name,
            format!(
                "f{}_k{}_d{}",
                case.frag_count, case.keys_per_frag, case.duplicate_every
            ),
        );
        group.bench_with_input(bench_id, &entries_list, |b, entries| {
            b.iter(|| {
                let (retained, obsolete) = merge_row_entries_for_bench(black_box(entries)).unwrap();
                black_box((retained, obsolete));
            });
        });
    }
    group.finish();
}

criterion_group! {
    name = benches;
    config = Criterion::default().sample_size(10);
    targets = bench_dedup_merge
}
criterion_main!(benches);