use std::{collections::BTreeMap, env, path::Path};

use criterion::{Criterion, criterion_group, criterion_main};
use tantivy::{
    Index, TantivyDocument,
    collector::DocSetCollector,
    query::AllQuery,
    schema::{Schema, Value},
};
use tempfile::TempDir;
use tici_grpc::CompactType;
use tici_shard::{
    compactor::shard_compactor::{ShardCompactor, ShardCompactorStat},
    reader::frag_reader::FragReaderMeta,
    test_utils::{
        TEST_INDEX_META_NORMAL, create_test_frag_path, test_gen_raw_docs, test_get_inf_range,
    },
    utils,
    utils::{
        key::ShardingKey,
        schema::{HANDLE_STORE_NAME, IndexMeta, VERSION_STORE_NAME},
        tantivy_doc::{RawDoc, TantivyDoc},
    },
    writer::frag_writer::FragWriterBuilder,
};
use tici_system::readable::ReadableSize;

struct FragData {
    tmp_dir: TempDir,
    frag_reader_metas: Vec<FragReaderMeta>,
    expected_raw_docs: BTreeMap<ShardingKey, RawDoc>,
    index_meta: IndexMeta,
}

fn setup_frag_data() -> FragData {
    let tmp_dir = TempDir::new().unwrap();
    let local_data_path = tmp_dir.path().to_str().unwrap();
    let index_meta = TEST_INDEX_META_NORMAL.clone();

    let (raw_docs, expected_raw_docs) = test_gen_raw_docs(5000000, 300000, 100000);

    let frag_count = 10;
    assert!(raw_docs.len() % frag_count == 0);
    let doc_count_per_frag = raw_docs.len() / frag_count;
    let mut frag_reader_metas = Vec::with_capacity(frag_count);
    for i in 0..frag_count {
        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 doc in &raw_docs[i * doc_count_per_frag..(i + 1) * doc_count_per_frag] {
            let tantivy_doc = TantivyDoc::new(doc, &index_meta).unwrap();
            frag_writer.add_tantivy_doc(tantivy_doc);
        }
        let frag_res = frag_writer.finalize().unwrap();

        let frag_reader_meta = FragReaderMeta::new(
            &frag_writer.frag_path(),
            frag_res.handle_stat.min_handle,
            frag_res.handle_stat.max_handle,
        )
        .unwrap();
        frag_reader_metas.push(frag_reader_meta);
    }

    FragData {
        tmp_dir,
        frag_reader_metas,
        expected_raw_docs,
        index_meta,
    }
}

fn copy_dir_all(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> std::io::Result<()> {
    std::fs::create_dir_all(&dst)?;
    for entry in std::fs::read_dir(src)? {
        let entry = entry?;
        let file_type = entry.file_type()?;
        if file_type.is_dir() {
            copy_dir_all(entry.path(), dst.as_ref().join(entry.file_name()))?;
        } else {
            std::fs::copy(entry.path(), dst.as_ref().join(entry.file_name()))?;
        }
    }
    Ok(())
}

fn new_frag_reader_metas(new_path: &str, frag_data: &FragData) -> Vec<FragReaderMeta> {
    copy_dir_all(frag_data.tmp_dir.path(), new_path).unwrap();
    let mut new_frag_reader_metas = frag_data.frag_reader_metas.clone();
    for frag_reader_meta in &mut new_frag_reader_metas {
        frag_reader_meta.local_frag_path = frag_reader_meta
            .local_frag_path
            .replace(frag_data.tmp_dir.path().to_str().unwrap(), new_path);
    }
    return new_frag_reader_metas;
}

fn get_version(doc: &TantivyDocument, schema: &Schema) -> u64 {
    let field = schema.get_field(VERSION_STORE_NAME).unwrap();
    let value = doc.get_first(field).unwrap();
    value.as_u64().unwrap()
}

fn get_handle(doc: &TantivyDocument, schema: &Schema) -> ShardingKey {
    let field = schema.get_field(HANDLE_STORE_NAME).unwrap();
    let value = doc.get_first(field).unwrap();
    ShardingKey::from_raw(value.as_bytes().unwrap().to_vec())
}

fn get_handle_versions(frag_path: &str) -> Vec<(ShardingKey, u64)> {
    let index = Index::open_in_dir(utils::frag_index_path(frag_path)).unwrap();
    let schema = index.schema();

    let index_reader = index.reader().unwrap();
    let searcher = index_reader.searcher();

    let docs = searcher.search(&AllQuery, &DocSetCollector).unwrap();
    let mut handle_versions = Vec::new();
    for addr in docs {
        let doc: TantivyDocument = searcher.doc(addr).unwrap();
        let h = get_handle(&doc, &schema);
        let v = get_version(&doc, &schema);
        handle_versions.push((h, v));
    }
    handle_versions
}

fn verify_compact_result(frag_path: &str, expected_docs: &BTreeMap<ShardingKey, RawDoc>) {
    let handle_versions = get_handle_versions(frag_path);
    assert_eq!(handle_versions.len(), expected_docs.len());
    for (h, v) in handle_versions {
        let raw_doc = expected_docs.get(&h).unwrap();
        assert_eq!(raw_doc.version, v);
    }
}

const BENCH_ITER_TIMES: usize = 10;

fn bench_compactor(c: &mut Criterion) {
    let frag_data = setup_frag_data();
    #[cfg(feature = "profiling")]
    let guard = pprof::ProfilerGuard::new(99).unwrap();

    let verify = match env::var("VERIFY") {
        Ok(v) => v,
        Err(_) => String::new(),
    };

    let tmp_dir = TempDir::new().unwrap();
    let test_path = tmp_dir.path().to_str().unwrap();
    let mut bench_iter_times = 0;
    let bench_name = "bench_compactor";
    c.bench_function(bench_name, |b| {
        b.iter(|| {
            let compact_path = format!("{}/{}", test_path, bench_iter_times);
            bench_iter_times += 1;
            let frag_reader_metas = new_frag_reader_metas(&compact_path, &frag_data);
            let output_path = format!("{}/output", compact_path);
            let compactor = ShardCompactor::new(
                frag_data.index_meta.clone(),
                output_path.clone(),
                frag_reader_metas,
                test_get_inf_range(),
                CompactType::Delta,
                1,
                128 * 1024 * 1024,
            );
            let rt = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .unwrap();
            let mut compact_stat = ShardCompactorStat::default();
            let (_new_seg_id, _handle_stat) =
                rt.block_on(async { compactor.compact(&mut compact_stat).await.unwrap() });
            println!("compact_stat={:?}", compact_stat);
            if !verify.is_empty() {
                verify_compact_result(&output_path, &frag_data.expected_raw_docs);
            }
        })
    });

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