use std::sync::Arc;

use criterion::{Criterion, criterion_group, criterion_main};
use log::{debug, info};
use tempfile::TempDir;
use tici_common::codec::number::NumberCodec;
use tici_shard::{
    reader::{
        frag_reader::FragReaderMeta,
        query_info::{QueryInfo, TantivySchemaExt, TopKRequest},
        shard_mirror::ShardMirror,
        shard_reader::ShardReader,
    },
    test_utils::{
        TEST_INDEX_META_NORMAL, TEST_MULTI_PK_META, create_test_frag_path, test_gen_raw_docs,
        test_gen_raw_docs_numeric_only, test_get_inf_range, test_get_shard_mirror_meta,
        test_get_shard_writer_meta,
    },
    utils::{
        schema::{ColumnInfoExt, SortField},
        tantivy_doc::TantivyDoc,
    },
    writer::frag_writer::FragWriterBuilder,
};

struct ShardReaderInfo {
    _tmp_dir: TempDir,
    reader: ShardReader,
    query_field: String,
    total_count: usize,
}

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 (raw_docs, _expected_raw_docs) = test_gen_raw_docs(1000000, 300000, 0);
    let frag_count = 10;
    let docs_per_frag = 130000;
    let mut frag_reader_metas = Vec::with_capacity(frag_count);
    for i in 0..frag_count {
        let raw_docs = raw_docs[i * docs_per_frag..(i + 1) * docs_per_frag].to_vec();
        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(tici_system::readable::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 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(
        test_get_shard_mirror_meta(&shard_writer_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,
    }
}

async fn setup_shard_reader_numeric_only() -> 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_MULTI_PK_META.clone();
    let schema = index_meta.get_tantivy_schema().clone();
    let raw_docs = test_gen_raw_docs_numeric_only(1_000_000);
    let frag_count = 10;
    let docs_per_frag = 100_000;
    let mut frag_reader_metas = Vec::with_capacity(frag_count);
    for i in 0..frag_count {
        let start_idx = i * docs_per_frag;
        let end_idx = std::cmp::min((i + 1) * docs_per_frag, raw_docs.len());
        let raw_docs_batch = raw_docs[start_idx..end_idx].to_vec();
        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(schema.clone())
            .index_num_threads(1)
            .index_mem_budget(tici_system::readable::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_batch {
            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 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,
    }
}

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

    let field_name = shard_reader_info.query_field.clone();
    let query_info = Arc::new(
        QueryInfo::query_word(
            &TEST_INDEX_META_NORMAL,
            &shard_reader_info.reader.get_schema().unwrap(),
            &[test_get_inf_range()],
            field_name,
            query.to_string(),
        )
        .unwrap(),
    );
    debug!("query_info: {:?}", query_info);
    let mut bench_iter_times = 0;
    c.bench_function(bench_name, |b| {
        b.iter(|| {
            bench_iter_times += 1;
            let (doc, stat) = shard_reader_info
                .reader
                .search_with_query_info(query_info.clone())
                .unwrap();
            debug!(
                "{}/{}/{} => {}/{}, {:?}",
                bench_iter_times,
                bench_name,
                query,
                doc.len(),
                shard_reader_info.total_count,
                stat
            );
        })
    });

    #[cfg(feature = "profiling")]
    if bench_iter_times > 0 {
        if let Ok(report) = guard.report().build() {
            let fname = format!(
                "bench_shard_reader_term_{}_{}.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_term_match_return_str(c: &mut Criterion) {
    let benches = vec![
        BenchInfo {
            name: "term_match_000100",
            query: "rwp00001",
        },
        BenchInfo {
            name: "term_match_001000",
            query: "rwp00010",
        },
        BenchInfo {
            name: "term_match_010000",
            query: "rwp00100",
        },
        BenchInfo {
            name: "term_match_100000",
            query: "rwp01000",
        },
        BenchInfo {
            name: "term_match_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(c, bench.name, bench.query, shard_reader_info.clone());
    }
}

fn run_bench_return_multi_int(c: &mut Criterion) {
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    let shard_reader_info = Arc::new(rt.block_on(setup_shard_reader_numeric_only()));

    let test_cases = vec![
        (1000000 as i64, "range_match_1000000"),
        (800000 as i64, "range_match_800000"),
        (10000 as i64, "range_match_10000"),
        (0 as i64, "range_match_0"),
    ];

    for (value, bench_name) in test_cases {
        #[cfg(feature = "profiling")]
        let guard = pprof::ProfilerGuard::new(99).unwrap();

        let local_shard_reader_info = shard_reader_info.clone();

        c.bench_function(bench_name, |b| {
            b.iter(|| {
                let read_ranges = vec![test_get_inf_range()];
                let schema = local_shard_reader_info.reader.get_schema().unwrap();
                let index_meta = TEST_MULTI_PK_META.clone();

                let mut ge_expr: tipb::Expr = tipb::Expr::default();
                ge_expr.set_tp(tipb::ExprType::ScalarFunc);
                ge_expr.set_sig(tipb::ScalarFuncSig::LeInt);

                let mut col_expr = tipb::Expr::default();
                col_expr.set_tp(tipb::ExprType::ColumnRef);
                col_expr.set_val("column_6".as_bytes().to_vec());
                ge_expr.mut_children().push(col_expr);

                let mut val_expr = tipb::Expr::default();
                val_expr.set_tp(tipb::ExprType::Int64);
                let mut buf = [0u8; 8];
                NumberCodec::encode_i64(&mut buf, value);
                val_expr.set_val(buf.to_vec());
                ge_expr.mut_children().push(val_expr);

                let return_field_names = index_meta
                    .pk_columns()
                    .iter()
                    .map(|pk_col| pk_col.store_name())
                    .collect::<Vec<String>>();

                let query_info = QueryInfo::parse_expr(
                    &index_meta,
                    &schema,
                    &ge_expr,
                    &read_ranges,
                    &return_field_names,
                    TopKRequest {
                        limit: 0,
                        sort_field_names: Vec::new(),
                        is_asc: Vec::new(),
                    },
                )
                .unwrap();

                local_shard_reader_info
                    .reader
                    .search_with_query_info(query_info)
                    .unwrap();
            })
        });
        #[cfg(feature = "profiling")]
        if let Ok(report) = guard.report().build() {
            let fname = format!(
                "bench_shard_reader_range_{}_{}.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();
        }
    }
}

fn run_bench_order_by_limit(c: &mut Criterion) {
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    let shard_reader_info = Arc::new(rt.block_on(setup_shard_reader_numeric_only()));

    let test_cases = vec![
        (10, "order_by_column4_limit_10"),
        (100, "order_by_column4_limit_100"),
        (1000, "order_by_column4_limit_1000"),
        (10000, "order_by_column4_limit_10000"),
    ];

    for (limit, bench_name) in test_cases {
        #[cfg(feature = "profiling")]
        let guard = pprof::ProfilerGuard::new(99).unwrap();

        let local_shard_reader_info = shard_reader_info.clone();
        let index_meta = TEST_MULTI_PK_META.clone();

        c.bench_function(bench_name, |b| {
            b.iter(|| {
                let read_ranges = vec![test_get_inf_range()];
                let schema = local_shard_reader_info.reader.get_schema().unwrap();

                // Create a query that returns all documents (no filtering)
                let query_info =
                    QueryInfo::query_all_for_test(&index_meta, &schema, read_ranges.clone())
                        .unwrap();

                let id_field = schema.doc_field("column_4").unwrap();
                let id_sf = SortField {
                    field: id_field,
                    is_asc: false,
                };
                let sort_fields = vec![id_sf];
                let query_info_with_sort = query_info
                    .add_topk(&schema, limit, sort_fields.clone(), false)
                    .unwrap();

                local_shard_reader_info
                    .reader
                    .search_with_query_info(query_info_with_sort)
                    .unwrap();
            })
        });

        #[cfg(feature = "profiling")]
        if let Ok(report) = guard.report().build() {
            let fname = format!(
                "bench_order_by_column4_{}_{}.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();
        }
    }
}

criterion_group! {
    name = benches;
    config = Criterion::default().sample_size(10);
    targets = run_bench_order_by_limit, run_bench_return_multi_int, run_bench_term_match_return_str,
}

criterion_main!(benches);