use std::{
    fs,
    path::{Path, PathBuf},
    time::{Duration, Instant},
};

use anyhow::{Context, Result, ensure};
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use tantivy::{Index, TantivyDocument};
use tempfile::TempDir;
use tici_shard::{
    compactor::frag_utils::{dedup_frags, merge_row_entries_for_bench},
    reader::frag_reader::FragReaderMeta,
    utils::{
        self, create_index_writer, key::ShardingKeyRange, row_entry::RowEntry,
        schema::HANDLE_STORE_NAME,
    },
};

#[derive(Clone, Copy)]
struct BenchCase {
    name: &'static str,
    frag_count: usize,
    keys_per_frag: usize,
    duplicate_every: usize,
}

struct BenchFixture {
    _tmp_dir: TempDir,
    source_root: PathBuf,
    source_frag_reader_metas: Vec<FragReaderMeta>,
    shard_range: ShardingKeyRange,
    input_key_count: usize,
    retained_key_count: usize,
    obsolete_key_count: usize,
}

struct BenchRun {
    run_root: PathBuf,
    obsolete_dir: PathBuf,
    frag_reader_metas: Vec<FragReaderMeta>,
}

fn build_entries(case: BenchCase) -> Vec<Vec<RowEntry>> {
    let mut entries_list = Vec::with_capacity(case.frag_count);

    for frag_idx in 0..case.frag_count {
        let mut entries = Vec::with_capacity(case.keys_per_frag);
        for key_idx in 0..case.keys_per_frag {
            let base_key_id = key_idx * (case.frag_count + 1);
            let key_id = if case.duplicate_every > 0 && key_idx % case.duplicate_every == 0 {
                base_key_id
            } else {
                base_key_id + frag_idx + 1
            };

            let key = format!("k{:020}", key_id).into_bytes();
            let del_mark = (frag_idx + key_idx) % 11 == 0;
            entries.push(RowEntry::new(key, del_mark, 0));
        }
        entries_list.push(entries);
    }

    entries_list
}

fn create_bench_frag(base_dir: &Path, name: &str, entries: &[RowEntry]) -> Result<FragReaderMeta> {
    let frag_dir = base_dir.join(name);
    fs::create_dir_all(&frag_dir)?;
    let frag_path = frag_dir
        .to_str()
        .context("frag_dir to_str failed")?
        .to_string();

    let entries_path = utils::frag_row_entry_path(&frag_path);
    RowEntry::write_all(&entries_path, entries)?;

    let index_path = utils::frag_index_path(&frag_path);
    fs::create_dir_all(&index_path)?;

    let mut schema_builder = tantivy::schema::Schema::builder();
    let handle_field = schema_builder.add_bytes_field(
        HANDLE_STORE_NAME,
        tantivy::schema::INDEXED | tantivy::schema::STORED,
    );
    let schema = schema_builder.build();
    let index = Index::create_in_dir(&index_path, schema)?;
    let mut index_writer = create_index_writer(&index, 1, 16 * 1024 * 1024, &frag_path)?;
    for entry in entries {
        let mut doc = TantivyDocument::new();
        doc.add_bytes(handle_field, &entry.key);
        index_writer.add_document(doc)?;
    }
    index_writer.commit()?;
    index_writer.wait_merging_threads()?;

    let min_handle = entries
        .first()
        .map(|entry| entry.key.clone())
        .unwrap_or_default();
    let max_handle = entries
        .last()
        .map(|entry| entry.key.clone())
        .unwrap_or_default();
    FragReaderMeta::new(&frag_path, min_handle, max_handle)
}

fn copy_dir_all(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> std::io::Result<()> {
    fs::create_dir_all(dst.as_ref())?;
    for entry in 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 {
            fs::copy(entry.path(), dst.as_ref().join(entry.file_name()))?;
        }
    }
    Ok(())
}

fn clone_frag_reader_metas(
    source_root: &Path,
    source_frag_reader_metas: &[FragReaderMeta],
    target_root: &Path,
) -> Result<Vec<FragReaderMeta>> {
    copy_dir_all(source_root, target_root)?;

    let source_root = source_root
        .to_str()
        .context("source_root to_str failed")?
        .to_string();
    let target_root = target_root
        .to_str()
        .context("target_root to_str failed")?
        .to_string();

    let mut frag_reader_metas = source_frag_reader_metas.to_vec();
    for frag_reader_meta in &mut frag_reader_metas {
        frag_reader_meta.local_frag_path =
            frag_reader_meta
                .local_frag_path
                .replacen(&source_root, &target_root, 1);
    }
    Ok(frag_reader_metas)
}

fn setup_fixture(case: BenchCase) -> Result<BenchFixture> {
    let tmp_dir = TempDir::new()?;
    let source_root = tmp_dir.path().join("source");
    fs::create_dir_all(&source_root)?;

    let entries_list = build_entries(case);
    let input_key_count: usize = entries_list.iter().map(Vec::len).sum();
    let (retained_key_count, obsolete_key_count) = merge_row_entries_for_bench(&entries_list)?;

    let mut source_frag_reader_metas = Vec::with_capacity(case.frag_count);
    for (frag_idx, entries) in entries_list.iter().enumerate() {
        source_frag_reader_metas.push(create_bench_frag(
            &source_root,
            &format!("f_{frag_idx}"),
            entries,
        )?);
    }

    println!(
        "case={}, frag_count={}, keys_per_frag={}, duplicate_every={}, input_key_count={}, retained_key_count={}, obsolete_key_count={}",
        case.name,
        case.frag_count,
        case.keys_per_frag,
        case.duplicate_every,
        input_key_count,
        retained_key_count,
        obsolete_key_count,
    );

    Ok(BenchFixture {
        _tmp_dir: tmp_dir,
        source_root,
        source_frag_reader_metas,
        shard_range: ShardingKeyRange::half_open(b"k".to_vec(), b"l".to_vec())?,
        input_key_count,
        retained_key_count,
        obsolete_key_count,
    })
}

fn prepare_run(fixture: &BenchFixture, case_name: &str, iter: u64) -> Result<BenchRun> {
    let run_root = fixture
        .source_root
        .parent()
        .context("fixture source root has no parent")?
        .join("runs")
        .join(format!("{case_name}_dedup_{iter}"));
    let obsolete_dir = run_root.join("obsolete");
    let frag_root = run_root.join("frags");

    let frag_reader_metas = clone_frag_reader_metas(
        &fixture.source_root,
        &fixture.source_frag_reader_metas,
        &frag_root,
    )?;
    fs::create_dir_all(&obsolete_dir)?;

    Ok(BenchRun {
        run_root,
        obsolete_dir,
        frag_reader_metas,
    })
}

fn run_dedup_once(fixture: &BenchFixture, case: BenchCase, bench_run: &BenchRun) -> Result<()> {
    let mut retained_count = 0usize;
    let output = dedup_frags(
        &fixture.shard_range,
        &bench_run.frag_reader_metas,
        |_entry| {
            retained_count += 1;
            Ok(())
        },
        &bench_run.obsolete_dir,
    )?;

    ensure!(
        retained_count == fixture.retained_key_count,
        "retained count mismatch, case={}, expected={}, actual={}",
        case.name,
        fixture.retained_key_count,
        retained_count,
    );
    ensure!(
        output.dedup_stat.input_key_count == fixture.input_key_count,
        "input key count mismatch, case={}, expected={}, actual={}",
        case.name,
        fixture.input_key_count,
        output.dedup_stat.input_key_count,
    );
    ensure!(
        output.dedup_stat.retained_key_count == fixture.retained_key_count,
        "dedup stat retained count mismatch, case={}, expected={}, actual={}",
        case.name,
        fixture.retained_key_count,
        output.dedup_stat.retained_key_count,
    );
    ensure!(
        output.dedup_stat.obsolete_by_dedup_count == fixture.obsolete_key_count,
        "obsolete key count mismatch, case={}, expected={}, actual={}",
        case.name,
        fixture.obsolete_key_count,
        output.dedup_stat.obsolete_by_dedup_count,
    );

    ensure!(
        output.obsolete_doc_bitsets.len() == fixture.source_frag_reader_metas.len(),
        "dedup should return one doc bitset per frag",
    );
    Ok(())
}

fn measure_run(fixture: &BenchFixture, case: BenchCase, iter: u64) -> Result<Duration> {
    let bench_run = prepare_run(fixture, case.name, iter)?;
    let start = Instant::now();
    let run_result = run_dedup_once(fixture, case, &bench_run);
    let elapsed = start.elapsed();
    let cleanup_result = fs::remove_dir_all(&bench_run.run_root).with_context(|| {
        format!(
            "remove bench run dir failed, case={}, run_root={}",
            case.name,
            bench_run.run_root.display()
        )
    });

    run_result?;
    cleanup_result?;
    Ok(elapsed)
}

fn bench_dedup_obsolete_modes(c: &mut Criterion) {
    let cases = [
        BenchCase {
            name: "medium_dup",
            frag_count: 32,
            keys_per_frag: 40_960,
            duplicate_every: 32,
        },
        BenchCase {
            name: "dense_dup",
            frag_count: 32,
            keys_per_frag: 40_960,
            duplicate_every: 8,
        },
        BenchCase {
            name: "high_frag_dense_dup",
            frag_count: 64,
            keys_per_frag: 20_480,
            duplicate_every: 8,
        },
    ];

    let mut group = c.benchmark_group("dedup_obsolete_doc_bitsets");
    for case in cases {
        let fixture = setup_fixture(case).unwrap();

        let bench_id = BenchmarkId::new(
            "build_doc_bitsets",
            format!(
                "{}_f{}_k{}_d{}",
                case.name, case.frag_count, case.keys_per_frag, case.duplicate_every
            ),
        );
        group.bench_function(bench_id, |b| {
            b.iter_custom(|iters| {
                let mut elapsed = Duration::ZERO;
                for iter in 0..iters {
                    elapsed += measure_run(&fixture, case, iter).unwrap();
                }
                elapsed
            });
        });
    }
    group.finish();
}

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