use std::{env, sync::Arc};

use criterion::{Criterion, criterion_group, criterion_main};
use log::{debug, info};
use tempfile::TempDir;
use tici_shard::{
    reader::{
        frag_reader::FragReaderMeta, query_info::QueryInfo, shard_mirror::ShardMirror,
        shard_reader::ShardReader,
    },
    test_utils::{
        TEST_INDEX_META_NORMAL, create_test_frag_path, test_gen_raw_docs, test_get_inf_range,
        test_get_shard_mirror_meta, test_get_shard_writer_meta,
    },
    utils::{key::ShardingKeyRange, schema::ColumnInfoExt, tantivy_doc::TantivyDoc},
    writer::frag_writer::FragWriterBuilder,
};
use tici_system::readable::ReadableSize;

struct ShardReaderInfo {
    _tmp_dir: TempDir,
    reader: ShardReader,
    query_field: String,
    total_count: usize,
    read_ranges: Vec<ShardingKeyRange>,
}

async fn setup_shard_reader() -> ShardReaderInfo {
    let tmp_dir = TempDir::new().unwrap();
    let local_data_path = tmp_dir.path().to_str().unwrap();
    info!("local_data_path={}", local_data_path);
    let index_meta = TEST_INDEX_META_NORMAL.clone();
    let frag_count = 10;
    let mut frag_reader_metas = Vec::with_capacity(frag_count);
    let mut read_ranges = Vec::with_capacity(frag_count);
    for i in 0..frag_count {
        let (raw_docs, _expected_raw_docs) = test_gen_raw_docs(100000, 0, 0);
        let frag_path = create_test_frag_path(tmp_dir.path()).unwrap();
        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(ReadableSize::mb(512).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);
        }
        let frag_res = frag_writer.finalize().unwrap();

        frag_reader_metas.push(
            FragReaderMeta::new(
                &frag_writer.frag_path(),
                frag_res.handle_stat.min_handle.clone(),
                frag_res.handle_stat.max_handle.clone(),
            )
            .unwrap(),
        );
        let min = frag_reader_metas[i].min_handle.raw();
        let max = frag_reader_metas[i].max_handle.raw();
        read_ranges.push(ShardingKeyRange::half_open(min.to_vec(), max.to_vec()).unwrap());
    }
    let total_doc_count = 100000 * frag_count;
    let shard_writer_meta = test_get_shard_writer_meta();
    let shard_mirror_meta = test_get_shard_mirror_meta(&shard_writer_meta);
    let dedup_output = ShardMirror::dedup_frags(
        shard_mirror_meta.shard_tracing_id(),
        shard_mirror_meta.shard_range().unwrap(),
        frag_reader_metas.clone(),
        format!("{}/temp", local_data_path),
    )
    .await
    .unwrap()
    .unwrap();
    let shard_mirror = ShardMirror::new(
        shard_mirror_meta,
        index_meta.clone(),
        frag_reader_metas,
        dedup_output.obsolete_doc_bitsets,
    )
    .await
    .unwrap();

    ShardReaderInfo {
        _tmp_dir: tmp_dir,
        reader: shard_mirror.get_reader(),
        query_field: index_meta.index_columns().first().unwrap().store_name(),
        total_count: total_doc_count,
        read_ranges,
    }
}

fn bench_shard_reader_count(
    c: &mut Criterion,
    bench_name: &str,
    query: &str,
    shard_reader_info: Arc<ShardReaderInfo>,
    read_ranges: &[ShardingKeyRange],
) {
    #[cfg(feature = "profiling")]
    let guard = pprof::ProfilerGuard::new(99).unwrap();

    let verify = env::var("VERIFY").unwrap_or_else(|_| "0".to_string()) == "1";
    let pushdown = env::var("PUSHDOWN").unwrap_or_else(|_| "1".to_string()) == "1";
    let inf_read_range = env::var("INF_READ_RANGE").unwrap_or_else(|_| "0".to_string()) == "1";
    let read_ranges = if inf_read_range {
        &[test_get_inf_range()]
    } else {
        read_ranges
    };
    info!(
        "VERIFY={}, PUSHDOWN={}, INF_READ_RANGE={}",
        verify, pushdown, inf_read_range
    );

    let query_count = Arc::new(
        QueryInfo::query_word_count(
            &TEST_INDEX_META_NORMAL,
            &shard_reader_info.reader.get_schema().unwrap(),
            read_ranges,
            shard_reader_info.query_field.clone(),
            query.to_string(),
        )
        .unwrap(),
    );
    let query_rows = Arc::new(
        QueryInfo::query_word(
            &TEST_INDEX_META_NORMAL,
            &shard_reader_info.reader.get_schema().unwrap(),
            read_ranges,
            shard_reader_info.query_field.clone(),
            query.to_string(),
        )
        .unwrap(),
    );
    let mut bench_iter_times = 0;
    c.bench_function(bench_name, |b| {
        b.iter(|| {
            bench_iter_times += 1;
            let (count1, stat) = if pushdown {
                shard_reader_info
                    .reader
                    .count_with_query_info(query_count.clone())
                    .unwrap()
            } else {
                let (rows_result, stat) = shard_reader_info
                    .reader
                    .search_with_query_info(query_rows.clone())
                    .unwrap();
                (rows_result.len() as u64, stat)
            };
            debug!(
                "{}/{}/{} => {}/{}, {:?}",
                bench_iter_times, bench_name, query, count1, shard_reader_info.total_count, stat
            );
            if verify {
                let count2 = if pushdown {
                    let (rows_result, _stat) = shard_reader_info
                        .reader
                        .search_with_query_info(query_rows.clone())
                        .unwrap();
                    rows_result.len() as u64
                } else {
                    let (count, _stat) = shard_reader_info
                        .reader
                        .count_with_query_info(query_count.clone())
                        .unwrap();
                    count
                };
                assert_eq!(count1, count2);
            }
        })
    });

    #[cfg(feature = "profiling")]
    if bench_iter_times > 0 {
        if let Ok(report) = guard.report().build() {
            let fname = format!(
                "bench_shard_reader_{}_{}.svg",
                bench_name,
                chrono::Local::now().format("%Y%m%d%H%M%S")
            );
            let file = std::fs::File::create(&fname).unwrap();
            report.flamegraph(file).unwrap();
        }
    }
}

struct BenchInfo {
    name: &'static str,
    query: &'static str,
}
fn run_bench(c: &mut Criterion) {
    let benches = vec![
        BenchInfo {
            name: "match_count_000100",
            query: "rwp00001",
        },
        BenchInfo {
            name: "match_count_001000",
            query: "rwp00010",
        },
        BenchInfo {
            name: "match_count_010000",
            query: "rwp00100",
        },
        BenchInfo {
            name: "match_count_100000",
            query: "rwp01000",
        },
        BenchInfo {
            name: "match_count_900000",
            query: "rwp09000",
        },
    ];

    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    let shard_reader_info = Arc::new(rt.block_on(setup_shard_reader()));

    for bench in benches.iter() {
        bench_shard_reader_count(
            c,
            bench.name,
            bench.query,
            shard_reader_info.clone(),
            &shard_reader_info.read_ranges,
        );
    }
}

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

criterion_main!(benches);