use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use flate2::bufread::GzDecoder;
use passgen::markov::{frequency::*, prefix::*, Markov};
use std::fs::File;
use std::io::{prelude::*, BufReader};

fn rockyou1m() -> Vec<String> {
    let file = File::open("benches/rockyou1m.txt.gz").unwrap();
    let reader = BufReader::new(file);
    let decoder = GzDecoder::new(reader);
    let reader = BufReader::new(decoder);
    reader.lines().map(|word| word.unwrap()).collect()
}

pub fn markov_benchmark_prefix<P: Prefix>(c: &mut Criterion, words: &[String]) {
    let type_name = std::any::type_name::<P>();
    let type_name = type_name.replace("passgen::markov::frequency::", "");
    let type_name = type_name.replace("passgen::markov::prefix::", "");
    let mut group = c.benchmark_group(&format!("Markov<{type_name}>"));
    group.sample_size(10);

    for rank in 0usize..5 {
        let benchmark = BenchmarkId::new("load", format!("rank{rank}"));
        group.bench_with_input(benchmark, &rank, |b, rank| {
            b.iter_with_large_drop(|| {
                let mut markov = Markov::<P>::new(*rank);
                for word in words.iter() {
                    markov.insert(word, 1);
                }
                markov
            })
        });
    }

    group.finish();
}

pub fn markov_benchmark(c: &mut Criterion) {
    let words = rockyou1m();

    // test frequency implementations
    markov_benchmark_prefix::<PrefixDefault<FrequencyHashMap>>(c, &words);
    markov_benchmark_prefix::<PrefixDefault<FrequencyBTree>>(c, &words);
    #[cfg(feature = "hashbrown")]
    markov_benchmark_prefix::<PrefixDefault<FrequencyHashbrown>>(c, &words);

    // test prefix implementations
    markov_benchmark_prefix::<PrefixHashMap<FrequencyDefault>>(c, &words);
    markov_benchmark_prefix::<PrefixBTree<FrequencyDefault>>(c, &words);
    #[cfg(feature = "hashbrown")]
    markov_benchmark_prefix::<PrefixHashbrown<FrequencyDefault>>(c, &words);
}

criterion_group!(benches, markov_benchmark);
criterion_main!(benches);