use criterion::{black_box, criterion_group, criterion_main, Criterion};
use radixtree::RadixTree;
use std::path::PathBuf;
use tempfile::tempdir;

fn criterion_benchmark(c: &mut Criterion) {
    // Create a temporary directory for benchmarks
    let temp_dir = tempdir().expect("Failed to create temp directory");
    let db_path = temp_dir.path().to_str().unwrap();
    
    // Benchmark set operation
    c.bench_function("set", |b| {
        let mut tree = RadixTree::new(db_path, true).unwrap();
        let mut i = 0;
        b.iter(|| {
            let key = format!("benchmark_key_{}", i);
            let value = format!("benchmark_value_{}", i).into_bytes();
            tree.set(&key, value).unwrap();
            i += 1;
        });
    });
    
    // Setup tree with data for get/list/delete benchmarks
    let mut setup_tree = RadixTree::new(db_path, true).unwrap();
    for i in 0..1000 {
        let key = format!("benchmark_key_{}", i);
        let value = format!("benchmark_value_{}", i).into_bytes();
        setup_tree.set(&key, value).unwrap();
    }
    
    // Benchmark get operation
    c.bench_function("get", |b| {
        let mut tree = RadixTree::new(db_path, false).unwrap();
        let mut i = 0;
        b.iter(|| {
            let key = format!("benchmark_key_{}", i % 1000);
            let _value = tree.get(&key).unwrap();
            i += 1;
        });
    });
    
    // Benchmark list operation
    c.bench_function("list", |b| {
        let mut tree = RadixTree::new(db_path, false).unwrap();
        b.iter(|| {
            let _keys = tree.list("benchmark_key_1").unwrap();
        });
    });
    
    // Benchmark getall operation
    c.bench_function("getall", |b| {
        let mut tree = RadixTree::new(db_path, false).unwrap();
        b.iter(|| {
            let _values = tree.getall("benchmark_key_1").unwrap();
        });
    });
    
    // Benchmark update operation
    c.bench_function("update", |b| {
        let mut tree = RadixTree::new(db_path, false).unwrap();
        let mut i = 0;
        b.iter(|| {
            let key = format!("benchmark_key_{}", i % 1000);
            let new_value = format!("updated_value_{}", i).into_bytes();
            tree.update(&key, new_value).unwrap();
            i += 1;
        });
    });
    
    // Benchmark delete operation
    c.bench_function("delete", |b| {
        // Create a fresh tree for deletion benchmarks
        let delete_dir = tempdir().expect("Failed to create temp directory");
        let delete_path = delete_dir.path().to_str().unwrap();
        let mut tree = RadixTree::new(delete_path, true).unwrap();
        
        // Setup keys to delete
        for i in 0..1000 {
            let key = format!("delete_key_{}", i);
            let value = format!("delete_value_{}", i).into_bytes();
            tree.set(&key, value).unwrap();
        }
        
        let mut i = 0;
        b.iter(|| {
            let key = format!("delete_key_{}", i % 1000);
            // Only try to delete if it exists
            if tree.get(&key).is_ok() {
                tree.delete(&key).unwrap();
            }
            i += 1;
        });
    });
    
    // Benchmark prefix operations with varying tree sizes
    let mut group = c.benchmark_group("prefix_operations");
    
    for &size in &[100, 1000, 10000] {
        // Create a fresh tree for each size
        let size_dir = tempdir().expect("Failed to create temp directory");
        let size_path = size_dir.path().to_str().unwrap();
        let mut tree = RadixTree::new(size_path, true).unwrap();
        
        // Insert data with common prefixes
        for i in 0..size {
            let prefix = match i % 5 {
                0 => "user",
                1 => "post",
                2 => "comment",
                3 => "product",
                _ => "category",
            };
            let key = format!("{}_{}", prefix, i);
            let value = format!("value_{}", i).into_bytes();
            tree.set(&key, value).unwrap();
        }
        
        // Benchmark list operation for this size
        group.bench_function(format!("list_size_{}", size), |b| {
            b.iter(|| {
                for prefix in &["user", "post", "comment", "product", "category"] {
                    let _keys = tree.list(prefix).unwrap();
                }
            });
        });
        
        // Benchmark getall operation for this size
        group.bench_function(format!("getall_size_{}", size), |b| {
            b.iter(|| {
                for prefix in &["user", "post", "comment", "product", "category"] {
                    let _values = tree.getall(prefix).unwrap();
                }
            });
        });
    }
    
    group.finish();
}

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