//! Vector Storage Performance Benchmarks
//!
//! Benchmarks for ConcurrentHNSWIndex at different scales:
//! - 10K vectors
//! - 100K vectors
//! - 1M vectors
//!
//! Run with: cargo bench --bench vector_bench

use criterion::{black_box, criterion_group, criterion_main, Criterion};
use rand::Rng;
use std::sync::Arc;
use tokio::runtime::Runtime;

use aifastdb_core::{
    types::HNSWConfig,
    vector::concurrent::{ConcurrentHNSWIndex, ConcurrentHNSWConfig},
};

fn generate_random_vector(dim: usize) -> Vec<f32> {
    let mut rng = rand::thread_rng();
    (0..dim).map(|_| rng.gen::<f32>()).collect()
}

fn generate_random_vectors(count: usize, dim: usize) -> Vec<(String, Vec<f32>)> {
    (0..count)
        .map(|i| (format!("vec_{}", i), generate_random_vector(dim)))
        .collect()
}

fn create_index_for_size(size: usize) -> ConcurrentHNSWIndex {
    let hnsw_config = HNSWConfig::new(128)
        .with_m(16)  // m_max0 auto-set to m*2=32
        .with_ef_construction(200)
        .with_ef_search(50)
        .with_max_elements(size + 1000); // Add buffer

    let config = ConcurrentHNSWConfig::with_hnsw(hnsw_config)
        .with_shards(4); // Use 4 shards for better parallelism

    ConcurrentHNSWIndex::new(config)
}

fn bench_insert_10k(c: &mut Criterion) {
    let mut group = c.benchmark_group("vector_insert_10k");
    group.sample_size(10);

    let size = 10_000;
    let dim = 128;
    let vectors = generate_random_vectors(size, dim);

    group.bench_function("insert_10k", |b| {
        let rt = Runtime::new().unwrap();
        b.iter(|| {
            let index = create_index_for_size(size);
            rt.block_on(async {
                for (id, vector) in vectors.iter() {
                    let _ = index.insert(id.clone(), vector.clone()).await;
                }
                // Flush to ensure all writes complete
                let _ = index.flush().await;
            });
            black_box(index.len())
        });
    });

    group.finish();
}

fn bench_insert_100k(c: &mut Criterion) {
    let mut group = c.benchmark_group("vector_insert_100k");
    group.sample_size(10);

    let size = 100_000;
    let dim = 128;
    let vectors = generate_random_vectors(size, dim);

    group.bench_function("insert_100k", |b| {
        let rt = Runtime::new().unwrap();
        b.iter(|| {
            let index = create_index_for_size(size);
            rt.block_on(async {
                // Insert in batches for better performance
                for chunk in vectors.chunks(1000) {
                    let batch: Vec<(String, Vec<f32>)> = chunk
                        .iter()
                        .map(|(id, vec)| (id.clone(), vec.clone()))
                        .collect();
                    let _ = index.insert_batch(batch).await;
                }
                let _ = index.flush().await;
            });
            black_box(index.len())
        });
    });

    group.finish();
}

fn bench_insert_1m(c: &mut Criterion) {
    let mut group = c.benchmark_group("vector_insert_1m");
    group.sample_size(10); // Criterion 0.5.x requires minimum 10 samples

    let size = 1_000_000;
    let dim = 128;
    let vectors = generate_random_vectors(size, dim);

    group.bench_function("insert_1m", |b| {
        let rt = Runtime::new().unwrap();
        b.iter(|| {
            let index = create_index_for_size(size);
            rt.block_on(async {
                // Insert in larger batches for 1M
                for chunk in vectors.chunks(5000) {
                    let batch: Vec<(String, Vec<f32>)> = chunk
                        .iter()
                        .map(|(id, vec)| (id.clone(), vec.clone()))
                        .collect();
                    let _ = index.insert_batch(batch).await;
                }
                let _ = index.flush().await;
            });
            black_box(index.len())
        });
    });

    group.finish();
}

fn bench_search_10k(c: &mut Criterion) {
    let mut group = c.benchmark_group("vector_search_10k");
    group.sample_size(100);

    let size = 10_000;
    let dim = 128;
    let index = create_index_for_size(size);
    let vectors = generate_random_vectors(size, dim);

    // Build index
    let rt = Runtime::new().unwrap();
    rt.block_on(async {
        for (id, vector) in vectors.iter() {
            let _ = index.insert(id.clone(), vector.clone()).await;
        }
        let _ = index.flush().await;
    });

    let query = generate_random_vector(dim);

    group.bench_function("search_top10", |b| {
        b.iter(|| {
            let results = index.search(black_box(&query), 10);
            black_box(results.len())
        });
    });

    group.bench_function("search_top100", |b| {
        b.iter(|| {
            let results = index.search(black_box(&query), 100);
            black_box(results.len())
        });
    });

    group.finish();
}

fn bench_search_100k(c: &mut Criterion) {
    let mut group = c.benchmark_group("vector_search_100k");
    group.sample_size(50);

    let size = 100_000;
    let dim = 128;
    let index = create_index_for_size(size);
    let vectors = generate_random_vectors(size, dim);

    // Build index
    let rt = Runtime::new().unwrap();
    rt.block_on(async {
        for chunk in vectors.chunks(1000) {
            let batch: Vec<(String, Vec<f32>)> = chunk
                .iter()
                .map(|(id, vec)| (id.clone(), vec.clone()))
                .collect();
            let _ = index.insert_batch(batch).await;
        }
        let _ = index.flush().await;
    });

    let query = generate_random_vector(dim);

    group.bench_function("search_top10", |b| {
        b.iter(|| {
            let results = index.search(black_box(&query), 10);
            black_box(results.len())
        });
    });

    group.bench_function("search_top100", |b| {
        b.iter(|| {
            let results = index.search(black_box(&query), 100);
            black_box(results.len())
        });
    });

    group.finish();
}

fn bench_search_1m(c: &mut Criterion) {
    let mut group = c.benchmark_group("vector_search_1m");
    group.sample_size(20);

    let size = 1_000_000;
    let dim = 128;
    let index = create_index_for_size(size);
    let vectors = generate_random_vectors(size, dim);

    // Build index
    let rt = Runtime::new().unwrap();
    rt.block_on(async {
        for chunk in vectors.chunks(5000) {
            let batch: Vec<(String, Vec<f32>)> = chunk
                .iter()
                .map(|(id, vec)| (id.clone(), vec.clone()))
                .collect();
            let _ = index.insert_batch(batch).await;
        }
        let _ = index.flush().await;
    });

    let query = generate_random_vector(dim);

    group.bench_function("search_top10", |b| {
        b.iter(|| {
            let results = index.search(black_box(&query), 10);
            black_box(results.len())
        });
    });

    group.bench_function("search_top100", |b| {
        b.iter(|| {
            let results = index.search(black_box(&query), 100);
            black_box(results.len())
        });
    });

    group.finish();
}

fn bench_concurrent_operations(c: &mut Criterion) {
    let mut group = c.benchmark_group("vector_concurrent");
    group.sample_size(20);

    let size = 10_000;
    let dim = 128;
    let index = Arc::new(create_index_for_size(size));
    let vectors = generate_random_vectors(size, dim);

    // Pre-populate index
    let rt = Runtime::new().unwrap();
    rt.block_on(async {
        for (id, vector) in vectors.iter() {
            let _ = index.insert(id.clone(), vector.clone()).await;
        }
        let _ = index.flush().await;
    });

    let query = generate_random_vector(dim);

    group.bench_function("concurrent_search", |b| {
        b.iter(|| {
            use rayon::prelude::*;
            let results: Vec<_> = (0..10)
                .into_par_iter()
                .map(|_| index.search(&query, 10))
                .collect();
            black_box(results.len())
        });
    });

    group.finish();
}

fn bench_filtered_search(c: &mut Criterion) {
    let mut group = c.benchmark_group("vector_filtered_search");
    group.sample_size(50);

    let size = 10_000;
    let dim = 128;
    let index = create_index_for_size(size);
    let vectors = generate_random_vectors(size, dim);

    // Build index with metadata (simulated)
    let rt = Runtime::new().unwrap();
    rt.block_on(async {
        for (id, vector) in vectors.iter() {
            let _ = index.insert(id.clone(), vector.clone()).await;
        }
        let _ = index.flush().await;
    });

    let query = generate_random_vector(dim);

    // Benchmark search without filter
    group.bench_function("search_no_filter", |b| {
        b.iter(|| {
            let results = index.search(black_box(&query), 100);
            black_box(results.len())
        });
    });

    // Note: MetadataFilter would be applied post-search in real usage
    // This benchmark shows the base search performance
    group.finish();
}

criterion_group!(
    benches,
    bench_insert_10k,
    bench_insert_100k,
    bench_insert_1m,
    bench_search_10k,
    bench_search_100k,
    bench_search_1m,
    bench_concurrent_operations,
    bench_filtered_search,
);

criterion_main!(benches);
