use criterion::{black_box, criterion_group, criterion_main, Criterion};

use implemented_trees::avl_tree::AVLTreeImpl;
use implemented_trees::red_black_tree::RedBlackTreeImpl;
use implemented_trees::binary_tree_ops::BinaryTreeOps;

// We will be testing these tree sizes as per the instructions.
const TREE_SIZES: [usize; 5] = [10_000, 40_000, 70_000, 100_000, 130_000];

// The worst-case scenario involves inserting increasing values.
// We will then search for the lowest (tree_size/10) values.
fn generate_increasing_sequence(n: usize) -> Vec<usize> {
    (0..n).collect()
}

// Benchmark insertion for AVL
fn bench_avl_insert(c: &mut Criterion) {
    let mut group = c.benchmark_group("AVL_Insert");
    for &size in &TREE_SIZES {
        let values = generate_increasing_sequence(size);
        group.bench_with_input(
            format!("insert_{}", size),
            &size,
            |b, &_s| {
                b.iter(|| {
                    let mut tree = AVLTreeImpl::new();
                    for v in &values {
                        tree.insert(black_box(*v));
                    }
                })
            }
        );
    }
    group.finish();
}

// Benchmark search for AVL
fn bench_avl_search(c: &mut Criterion) {
    let mut group = c.benchmark_group("AVL_Search");
    for &size in &TREE_SIZES {
        let values = generate_increasing_sequence(size);
        let mut tree = AVLTreeImpl::new();
        // Pre-insert the tree outside of b.iter() to measure search only.
        for v in &values {
            tree.insert(*v);
        }

        let num_searches = size / 10;
        let search_values = &values[0..num_searches];

        group.bench_with_input(
            format!("search_{}", size),
            &size,
            |b, &_s| {
                b.iter(|| {
                    for &sv in search_values {
                        let _ = tree.find_node(&black_box(sv));
                    }
                })
            }
        );
    }
    group.finish();
}

// Benchmark insertion for Red-Black
fn bench_rb_insert(c: &mut Criterion) {
    let mut group = c.benchmark_group("RB_Insert");
    for &size in &TREE_SIZES {
        let values = generate_increasing_sequence(size);
        group.bench_with_input(
            format!("insert_{}", size),
            &size,
            |b, &_s| {
                b.iter(|| {
                    let mut tree = RedBlackTreeImpl::new();
                    for v in &values {
                        tree.insert(black_box(*v));
                    }
                })
            }
        );
    }
    group.finish();
}

// Benchmark search for Red-Black
fn bench_rb_search(c: &mut Criterion) {
    let mut group = c.benchmark_group("RB_Search");
    for &size in &TREE_SIZES {
        let values = generate_increasing_sequence(size);
        let mut tree = RedBlackTreeImpl::new();
        // Pre-insert the tree.
        for v in &values {
            tree.insert(*v);
        }

        let num_searches = size / 10;
        let search_values = &values[0..num_searches];

        group.bench_with_input(
            format!("search_{}", size),
            &size,
            |b, &_s| {
                b.iter(|| {
                    for &sv in search_values {
                        let _ = tree.find_node(&black_box(sv));
                    }
                })
            }
        );
    }
    group.finish();
}

criterion_group!(benches, bench_avl_insert, bench_avl_search, bench_rb_insert, bench_rb_search);
criterion_main!(benches);