use std::{sync::Arc, time::Instant};

use anyhow::{Context, Result, anyhow};
use criterion::{Criterion, criterion_group, criterion_main};
use log::{debug, info};
use tempfile::TempDir;
use tici_search_lib::searcher::search_shard_reader;
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::{schema::ColumnInfoExt, 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() -> Result<ShardReaderInfo> {
    let tmp_dir =
        TempDir::new().context("failed to create temporary directory for shard reader")?;
    let local_data_path = tmp_dir
        .path()
        .to_str()
        .context("temporary directory path is not valid UTF-8")?;
    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())
            .with_context(|| format!("failed to create test frag path for frag {}", i))?;
        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()
            .with_context(|| format!("failed to build frag writer for frag {}", i))?;
        for raw_doc in &raw_docs {
            let tantivy_doc = TantivyDoc::new(raw_doc, &index_meta)
                .with_context(|| format!("failed to build tantivy doc for frag {}", i))?;
            frag_writer.add_tantivy_doc(tantivy_doc);
        }
        let frag_res = frag_writer
            .finalize()
            .with_context(|| format!("failed to finalize frag {}", i))?;

        frag_reader_metas.push(
            FragReaderMeta::new(
                &frag_writer.frag_path(),
                frag_res.handle_stat.min_handle.clone(),
                frag_res.handle_stat.max_handle.clone(),
            )
            .with_context(|| format!("failed to build FragReaderMeta for frag {}", i))?,
        );
    }
    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()
            .context("failed to get shard range")?,
        frag_reader_metas.clone(),
        format!("{}/temp", local_data_path),
    )
    .await
    .context("dedup frags task failed to join")?
    .context("dedup frags failed")?;
    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
    .context("failed to create shard mirror")?;

    let query_field = index_meta
        .index_columns()
        .first()
        .context("index meta has no index columns")?
        .store_name();

    Ok(ShardReaderInfo {
        _tmp_dir: tmp_dir,
        reader: shard_mirror.get_reader(),
        query_field,
        total_count: total_doc_count,
    })
}

fn search_once(
    shard_reader_info: &ShardReaderInfo,
    query_info: Arc<QueryInfo>,
    bench_name: &str,
    query: &str,
    bench_iter_times: u64,
) -> Result<()> {
    let (result, stat) = search_shard_reader(&shard_reader_info.reader, query_info)
        .map_err(|err| anyhow!("search_shard_reader failed: {err}"))?;
    debug!(
        "{}/{}/{} => {}/{}, {:?}",
        bench_iter_times, bench_name, query, result.row_count, shard_reader_info.total_count, stat
    );
    Ok(())
}

fn bench_search_shard_reader_rows(
    c: &mut Criterion,
    bench_name: &str,
    query: &str,
    shard_reader_info: Arc<ShardReaderInfo>,
) -> Result<()> {
    let field_name = shard_reader_info.query_field.clone();
    let schema = shard_reader_info
        .reader
        .get_schema()
        .context("failed to get shard reader schema")?;
    let query_info = Arc::new(
        QueryInfo::query_word(
            &TEST_INDEX_META_NORMAL,
            &schema,
            &[test_get_inf_range()],
            field_name,
            query.to_string(),
        )
        .with_context(|| format!("failed to build query info for query {}", query))?,
    );
    debug!("query_info: {:?}", query_info);

    let mut bench_iter_times = 0_u64;
    let mut bench_result = Ok(());
    c.bench_function(bench_name, |b| {
        b.iter_custom(|iters| {
            let watcher = Instant::now();
            for _ in 0..iters {
                bench_iter_times += 1;
                if let Err(err) = search_once(
                    &shard_reader_info,
                    query_info.clone(),
                    bench_name,
                    query,
                    bench_iter_times,
                ) {
                    bench_result = Err(err);
                    break;
                }
            }
            watcher.elapsed()
        })
    });
    bench_result
}

struct BenchInfo {
    name: &'static str,
    query: &'static str,
}

fn run_bench_term_match_return_rows(c: &mut Criterion) -> Result<()> {
    let benches = vec![
        BenchInfo {
            name: "searchlib_term_match_000100",
            query: "rwp00001",
        },
        BenchInfo {
            name: "searchlib_term_match_001000",
            query: "rwp00010",
        },
        BenchInfo {
            name: "searchlib_term_match_010000",
            query: "rwp00100",
        },
        BenchInfo {
            name: "searchlib_term_match_100000",
            query: "rwp01000",
        },
        BenchInfo {
            name: "searchlib_term_match_900000",
            query: "rwp09000",
        },
    ];

    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .context("failed to build tokio runtime for benchmark")?;
    let shard_reader_info = Arc::new(rt.block_on(setup_shard_reader())?);

    for bench in benches.iter() {
        bench_search_shard_reader_rows(c, bench.name, bench.query, shard_reader_info.clone())?;
    }
    Ok(())
}

fn run_bench(c: &mut Criterion) {
    run_bench_term_match_return_rows(c).unwrap();
}

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

criterion_main!(benches);