use criterion::{Criterion, criterion_group, criterion_main};
use tempfile::TempDir;
use tici_shard::{
    test_utils::{TEST_INDEX_META_OVERLAP, test_gen_raw_docs},
    utils::tantivy_doc::TantivyDoc,
    writer::frag_writer::FragWriterBuilder,
};
fn bench_frag_writer(c: &mut Criterion) {
    let tmp_dir = TempDir::new().unwrap();
    let local_data_path = tmp_dir.path().to_str().unwrap();
    let index_meta = TEST_INDEX_META_OVERLAP.clone();
    let (raw_docs, _expected_raw_docs) = test_gen_raw_docs(100000, 30000, 10000);
    let mut count = 0;
    c.bench_function("FragWriter add_tantivy_doc", |b| {
        b.iter(|| {
            let frag_path = format!("{}/frag_{}", local_data_path, count);
            count += 1;
            let mut frag_writer = FragWriterBuilder::new()
                .frag_base_path(local_data_path)
                .frag_path(&frag_path)
                .schema(index_meta.get_tantivy_schema().clone())
                .index_num_threads(1)
                .index_mem_budget(tici_system::readable::ReadableSize::mb(256).0 as usize)
                .is_base(false)
                .analyzer_manager(index_meta.get_analyzer_manager().clone())
                .sort_fields(vec![])
                .build()
                .unwrap();

            for raw_doc in &raw_docs {
                let tantivy_doc = TantivyDoc::new(raw_doc, &index_meta).unwrap();
                frag_writer.add_tantivy_doc(tantivy_doc);
            }
            frag_writer.finalize().unwrap();
        });
    });
}

criterion_group!(benches, bench_frag_writer);
criterion_main!(benches);