//! Index comparison benchmarks for Dakera
//!
//! Compares accuracy and speed tradeoffs between different index types:
//! - Flat (brute-force) search
//! - HNSW index
//! - IVF index

use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use std::sync::Arc;

use common::{DistanceMetric, QueryRequest, Vector};
use engine::{HnswConfig, HnswIndex, IvfConfig, IvfIndex, SearchEngine};
use storage::InMemoryStorage;
use storage::VectorStorage;

/// Generate random vectors for benchmarking
fn generate_vectors(count: usize, dimension: usize, prefix: &str) -> Vec<Vector> {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};

    (0..count)
        .map(|i| {
            let mut hasher = DefaultHasher::new();
            i.hash(&mut hasher);
            let seed = hasher.finish();

            let values: Vec<f32> = (0..dimension)
                .map(|j| {
                    let mut h = DefaultHasher::new();
                    (seed + j as u64).hash(&mut h);
                    (h.finish() as f32 / u64::MAX as f32) * 2.0 - 1.0
                })
                .collect();

            Vector {
                id: format!("{}{}", prefix, i),
                values,
                metadata: None,
                ttl_seconds: None,
                expires_at: None,
            }
        })
        .collect()
}

/// Generate a random query vector
fn generate_query_vector(dimension: usize, seed: u64) -> Vec<f32> {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};

    (0..dimension)
        .map(|j| {
            let mut h = DefaultHasher::new();
            (seed + j as u64).hash(&mut h);
            (h.finish() as f32 / u64::MAX as f32) * 2.0 - 1.0
        })
        .collect()
}

fn bench_flat_vs_hnsw(c: &mut Criterion) {
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();

    let dimension = 128;
    let top_k = 10;

    let mut group = c.benchmark_group("flat_vs_hnsw");
    group.throughput(Throughput::Elements(1));

    for dataset_size in [1_000, 5_000, 10_000].iter() {
        // Generate test data
        let vectors = generate_vectors(*dataset_size, dimension, "v");
        let query_vector = generate_query_vector(dimension, 42);

        // Benchmark flat search (brute force via SearchEngine)
        {
            let storage = Arc::new(InMemoryStorage::new());
            let engine = Arc::new(SearchEngine::new(storage.clone()));
            let namespace = "bench".to_string();

            rt.block_on(async {
                storage.ensure_namespace(&namespace).await.unwrap();
                storage.upsert(&namespace, vectors.clone()).await.unwrap();
            });

            group.bench_with_input(
                BenchmarkId::new("flat", dataset_size),
                dataset_size,
                |b, _| {
                    b.iter(|| {
                        rt.block_on(async {
                            let request = QueryRequest {
                                vector: query_vector.clone(),
                                top_k,
                                distance_metric: DistanceMetric::Cosine,
                                include_metadata: false,
                                include_vectors: false,
                                filter: None,
                                cursor: None,
                            };
                            let response = engine.search(&namespace, &request).await.unwrap();
                            black_box(response)
                        })
                    });
                },
            );
        }

        // Benchmark HNSW index
        {
            let config = HnswConfig::new(16, 100, 50).with_distance_metric(DistanceMetric::Cosine);
            let hnsw = HnswIndex::with_config(config);

            // Build index
            for v in &vectors {
                hnsw.insert(v.id.clone(), v.values.clone());
            }

            group.bench_with_input(
                BenchmarkId::new("hnsw", dataset_size),
                dataset_size,
                |b, _| {
                    b.iter(|| {
                        let results = hnsw.search(&query_vector, top_k);
                        black_box(results)
                    });
                },
            );
        }
    }
    group.finish();
}

fn bench_hnsw_parameters(c: &mut Criterion) {
    let dimension = 128;
    let dataset_size = 10_000;
    let top_k = 10;

    let vectors = generate_vectors(dataset_size, dimension, "v");
    let query_vector = generate_query_vector(dimension, 42);

    let mut group = c.benchmark_group("hnsw_ef_search");
    group.throughput(Throughput::Elements(1));

    // Test different ef_search values
    for ef_search in [10, 25, 50, 100, 200].iter() {
        let config =
            HnswConfig::new(16, 100, *ef_search).with_distance_metric(DistanceMetric::Cosine);
        let hnsw = HnswIndex::with_config(config);

        for v in &vectors {
            hnsw.insert(v.id.clone(), v.values.clone());
        }

        group.bench_with_input(BenchmarkId::from_parameter(ef_search), ef_search, |b, _| {
            b.iter(|| {
                let results = hnsw.search(&query_vector, top_k);
                black_box(results)
            });
        });
    }
    group.finish();
}

fn bench_ivf_parameters(c: &mut Criterion) {
    let dimension = 128;
    let dataset_size = 10_000;
    let top_k = 10;

    let vectors = generate_vectors(dataset_size, dimension, "v");
    let query_vector = generate_query_vector(dimension, 42);

    let mut group = c.benchmark_group("ivf_nprobe");
    group.throughput(Throughput::Elements(1));

    // Test different nprobe values
    for nprobe in [1, 2, 4, 8, 16].iter() {
        let config = IvfConfig {
            n_clusters: 100,
            n_probe: *nprobe,
            max_iterations: 100,
            convergence_threshold: 1e-4,
            metric: DistanceMetric::Cosine,
        };

        let mut ivf = IvfIndex::new(config);

        // Train on sample data
        let training_data: Vec<Vec<f32>> = vectors.iter().map(|v| v.values.clone()).collect();
        ivf.train(&training_data).unwrap();

        // Add vectors
        for v in &vectors {
            ivf.add(v.id.clone(), v.values.clone()).unwrap();
        }

        group.bench_with_input(BenchmarkId::from_parameter(nprobe), nprobe, |b, _| {
            b.iter(|| {
                let results = ivf.search(&query_vector, top_k);
                black_box(results)
            });
        });
    }
    group.finish();
}

fn bench_index_build_time(c: &mut Criterion) {
    let dimension = 128;

    let mut group = c.benchmark_group("index_build_time");

    for dataset_size in [1_000, 5_000, 10_000].iter() {
        group.throughput(Throughput::Elements(*dataset_size as u64));

        let vectors = generate_vectors(*dataset_size, dimension, "v");

        // HNSW build time
        group.bench_with_input(
            BenchmarkId::new("hnsw_build", dataset_size),
            dataset_size,
            |b, _| {
                b.iter(|| {
                    let config =
                        HnswConfig::new(16, 100, 50).with_distance_metric(DistanceMetric::Cosine);
                    let hnsw = HnswIndex::with_config(config);

                    for v in &vectors {
                        hnsw.insert(v.id.clone(), v.values.clone());
                    }
                    black_box(hnsw)
                });
            },
        );

        // IVF build time (including training)
        group.bench_with_input(
            BenchmarkId::new("ivf_build", dataset_size),
            dataset_size,
            |b, _| {
                b.iter(|| {
                    let config = IvfConfig {
                        n_clusters: 100,
                        n_probe: 4,
                        max_iterations: 100,
                        convergence_threshold: 1e-4,
                        metric: DistanceMetric::Cosine,
                    };
                    let mut ivf = IvfIndex::new(config);

                    let training_data: Vec<Vec<f32>> =
                        vectors.iter().map(|v| v.values.clone()).collect();
                    ivf.train(&training_data).unwrap();

                    for v in &vectors {
                        ivf.add(v.id.clone(), v.values.clone()).unwrap();
                    }
                    black_box(ivf)
                });
            },
        );
    }
    group.finish();
}

fn bench_memory_footprint(c: &mut Criterion) {
    let dimension = 128;
    let dataset_size = 10_000;

    let vectors = generate_vectors(dataset_size, dimension, "v");

    let mut group = c.benchmark_group("memory_operations");
    group.throughput(Throughput::Elements(1));

    // Measure memory allocation pattern through repeated operations
    group.bench_function("hnsw_insert_pattern", |b| {
        b.iter(|| {
            let config = HnswConfig::new(16, 100, 50).with_distance_metric(DistanceMetric::Cosine);
            let hnsw = HnswIndex::with_config(config);

            // Insert 1000 vectors
            for v in vectors.iter().take(1000) {
                hnsw.insert(v.id.clone(), v.values.clone());
            }
            black_box(hnsw.stats())
        });
    });

    group.finish();
}

criterion_group!(
    benches,
    bench_flat_vs_hnsw,
    bench_hnsw_parameters,
    bench_ivf_parameters,
    bench_index_build_time,
    bench_memory_footprint,
);

criterion_main!(benches);
