//! Lance vs JSONL Storage Performance Benchmarks
//!
//! Comprehensive benchmark suite comparing Lance columnar storage against
//! the legacy JSONL/mmap storage backend across multiple dimensions:
//!
//! - **Document CRUD**: Single/batch write, random read, sequential scan
//! - **Filtered Queries**: Tag, metadata, time-range, importance filters
//! - **Vector Search**: Lance ANN (IVF_PQ) vs HNSW recall & throughput
//! - **Storage Efficiency**: Disk space, memory usage
//! - **Version Control**: Time-travel overhead
//!
//! Run with:
//!   cargo bench --bench lance_bench --features lance-storage
//!
//! Run specific group:
//!   cargo bench --bench lance_bench --features lance-storage -- "doc_write"

use criterion::{
    black_box, criterion_group, criterion_main, Criterion, BenchmarkId, Throughput,
};
use rand::Rng;
use tempfile::TempDir;

use aifastdb_core::{
    storage::{
        backend::{StorageBackend, QueryFilter, MetadataFilter},
        DocumentStore,
    },
    types::{ContentType, MemoryNode},
};

#[cfg(feature = "lance-storage")]
use aifastdb_core::storage::lance::{
    LanceStorageBackend,
    lance_vector::{LanceVectorBackend, LanceVectorConfig},
};

#[cfg(feature = "lance-storage")]
use aifastdb_core::storage::lance::lance_backend::LanceConfig;

#[cfg(feature = "lance-storage")]
use aifastdb_core::storage::vector_backend::{
    VectorBackend, VectorIndexConfig, VectorIndexType, VectorDistanceMetric,
};

// ============================================================================
// Helper functions
// ============================================================================

/// Generate a random embedding vector of given dimension (f32).
fn random_embedding(dim: usize) -> Vec<f32> {
    let mut rng = rand::thread_rng();
    (0..dim).map(|_| rng.gen::<f32>()).collect()
}

/// Generate a MemoryNode with random content, embedding, tags, and metadata.
fn generate_node(id_suffix: &str, dim: usize) -> MemoryNode {
    let mut rng = rand::thread_rng();
    let mut node = MemoryNode::new(
        format!("Content for node {} — random data: {}", id_suffix, rng.gen::<u64>()),
    );
    node.id = format!("node-{}", id_suffix);
    node.embedding = Some(random_embedding(dim));
    node.tags = vec![
        format!("tag-{}", rng.gen_range(0..10u32)),
        format!("category-{}", rng.gen_range(0..5u32)),
    ];
    node.content_type = match rng.gen_range(0..3u32) {
        0 => ContentType::Text,
        1 => ContentType::Code,
        _ => ContentType::Conversation,
    };
    node.metadata.insert(
        "priority".to_string(),
        serde_json::Value::Number(serde_json::Number::from(rng.gen_range(1..=5i64))),
    );
    node.metadata.insert(
        "source".to_string(),
        serde_json::Value::String(format!("source-{}", rng.gen_range(0..20u32))),
    );
    node
}

/// Generate a batch of MemoryNodes.
fn generate_nodes(count: usize, dim: usize) -> Vec<MemoryNode> {
    (0..count)
        .map(|i| generate_node(&format!("{:06}", i), dim))
        .collect()
}

/// Create a JSONL-backed DocumentStore in a temp directory.
fn create_jsonl_store(dir: &TempDir) -> DocumentStore {
    let path = dir.path().join("documents.jsonl");
    DocumentStore::open_or_create(&path).unwrap()
}

/// Create a Lance-backed StorageBackend in a temp directory.
#[cfg(feature = "lance-storage")]
fn create_lance_store(dir: &TempDir, dim: usize) -> LanceStorageBackend {
    let path = dir.path().join("lance_store.lance");
    let config = LanceConfig::new(&path)
        .with_embedding_dim(dim as i32)
        .with_write_batch_size(1024);
    LanceStorageBackend::open(config).unwrap()
}

/// Pre-populate a StorageBackend with N nodes and flush.
fn populate_backend(store: &dyn StorageBackend, nodes: &[MemoryNode]) {
    for node in nodes {
        store.put(node.clone()).unwrap();
    }
    store.flush().unwrap();
}

// ============================================================================
// Benchmark Group 1: Document Write Performance
// ============================================================================

fn bench_doc_write_single(c: &mut Criterion) {
    let mut group = c.benchmark_group("doc_write_single");
    group.sample_size(10);
    let dim = 128;

    for count in [100, 1_000, 5_000] {
        let nodes = generate_nodes(count, dim);

        // JSONL
        group.throughput(Throughput::Elements(count as u64));
        group.bench_with_input(
            BenchmarkId::new("jsonl", count),
            &count,
            |b, &_count| {
                b.iter_with_setup(
                    || {
                        let dir = TempDir::new().unwrap();
                        let store = create_jsonl_store(&dir);
                        (dir, store)
                    },
                    |(_dir, store)| {
                        for node in &nodes {
                            store.put(node.clone()).unwrap();
                        }
                        store.flush().unwrap();
                        black_box(store.len())
                    },
                );
            },
        );

        // Lance
        #[cfg(feature = "lance-storage")]
        group.bench_with_input(
            BenchmarkId::new("lance", count),
            &count,
            |b, &_count| {
                b.iter_with_setup(
                    || {
                        let dir = TempDir::new().unwrap();
                        let store = create_lance_store(&dir, dim);
                        (dir, store)
                    },
                    |(_dir, store)| {
                        for node in &nodes {
                            store.put(node.clone()).unwrap();
                        }
                        store.flush().unwrap();
                        black_box(store.len().unwrap())
                    },
                );
            },
        );
    }

    group.finish();
}

fn bench_doc_write_batch(c: &mut Criterion) {
    let mut group = c.benchmark_group("doc_write_batch");
    group.sample_size(10);
    let dim = 128;

    for count in [1_000, 10_000] {
        let nodes = generate_nodes(count, dim);

        group.throughput(Throughput::Elements(count as u64));

        // JSONL (batch = put_batch via StorageBackend)
        group.bench_with_input(
            BenchmarkId::new("jsonl", count),
            &count,
            |b, &_count| {
                b.iter_with_setup(
                    || {
                        let dir = TempDir::new().unwrap();
                        let store = create_jsonl_store(&dir);
                        (dir, store)
                    },
                    |(_dir, store)| {
                        let backend: &dyn StorageBackend = &store;
                        backend.put_batch(nodes.clone()).unwrap();
                        backend.flush().unwrap();
                        black_box(backend.len().unwrap())
                    },
                );
            },
        );

        // Lance (batch = put_batch)
        #[cfg(feature = "lance-storage")]
        group.bench_with_input(
            BenchmarkId::new("lance", count),
            &count,
            |b, &_count| {
                b.iter_with_setup(
                    || {
                        let dir = TempDir::new().unwrap();
                        let store = create_lance_store(&dir, dim);
                        (dir, store)
                    },
                    |(_dir, store)| {
                        store.put_batch(nodes.clone()).unwrap();
                        store.flush().unwrap();
                        black_box(store.len().unwrap())
                    },
                );
            },
        );
    }

    group.finish();
}

// ============================================================================
// Benchmark Group 2: Document Read Performance
// ============================================================================

fn bench_doc_read_random(c: &mut Criterion) {
    let mut group = c.benchmark_group("doc_read_random");
    group.sample_size(50);
    let dim = 128;
    let count = 5_000;
    let nodes = generate_nodes(count, dim);

    // Collect IDs for random access
    let ids: Vec<String> = nodes.iter().map(|n| n.id.clone()).collect();

    // JSONL
    {
        let dir = TempDir::new().unwrap();
        let store = create_jsonl_store(&dir);
        populate_backend(&store, &nodes);

        group.throughput(Throughput::Elements(100));
        group.bench_function("jsonl/100_gets", |b| {
            let mut rng = rand::thread_rng();
            b.iter(|| {
                for _ in 0..100 {
                    let idx = rng.gen_range(0..ids.len());
                    let node = store.get(&ids[idx]);
                    black_box(node);
                }
            });
        });
    }

    // Lance
    #[cfg(feature = "lance-storage")]
    {
        let dir = TempDir::new().unwrap();
        let store = create_lance_store(&dir, dim);
        populate_backend(&store, &nodes);

        group.throughput(Throughput::Elements(100));
        group.bench_function("lance/100_gets", |b| {
            let mut rng = rand::thread_rng();
            b.iter(|| {
                for _ in 0..100 {
                    let idx = rng.gen_range(0..ids.len());
                    let node = store.get(&ids[idx]).unwrap();
                    black_box(node);
                }
            });
        });
    }

    group.finish();
}

fn bench_doc_scan_all(c: &mut Criterion) {
    let mut group = c.benchmark_group("doc_scan_all");
    group.sample_size(20);
    let dim = 128;

    for count in [1_000, 5_000, 10_000] {
        let nodes = generate_nodes(count, dim);

        // JSONL
        {
            let dir = TempDir::new().unwrap();
            let store = create_jsonl_store(&dir);
            populate_backend(&store, &nodes);

            group.throughput(Throughput::Elements(count as u64));
            group.bench_with_input(
                BenchmarkId::new("jsonl", count),
                &count,
                |b, _| {
                    b.iter(|| {
                        let all = store.all();
                        black_box(all.len())
                    });
                },
            );
        }

        // Lance
        #[cfg(feature = "lance-storage")]
        {
            let dir = TempDir::new().unwrap();
            let store = create_lance_store(&dir, dim);
            populate_backend(&store, &nodes);

            group.throughput(Throughput::Elements(count as u64));
            group.bench_with_input(
                BenchmarkId::new("lance", count),
                &count,
                |b, _| {
                    b.iter(|| {
                        let all = store.all().unwrap();
                        black_box(all.len())
                    });
                },
            );
        }
    }

    group.finish();
}

// ============================================================================
// Benchmark Group 3: Filtered Query Performance
// ============================================================================

fn bench_query_by_tag(c: &mut Criterion) {
    let mut group = c.benchmark_group("query_by_tag");
    group.sample_size(30);
    let dim = 128;
    let count = 10_000;
    let nodes = generate_nodes(count, dim);

    // JSONL
    {
        let dir = TempDir::new().unwrap();
        let store = create_jsonl_store(&dir);
        populate_backend(&store, &nodes);
        let backend: &dyn StorageBackend = &store;

        group.bench_function("jsonl", |b| {
            b.iter(|| {
                let results = backend.find_by_tag("tag-3").unwrap();
                black_box(results.len())
            });
        });
    }

    // Lance
    #[cfg(feature = "lance-storage")]
    {
        let dir = TempDir::new().unwrap();
        let store = create_lance_store(&dir, dim);
        populate_backend(&store, &nodes);

        group.bench_function("lance", |b| {
            b.iter(|| {
                let results = store.find_by_tag("tag-3").unwrap();
                black_box(results.len())
            });
        });
    }

    group.finish();
}

fn bench_query_by_metadata(c: &mut Criterion) {
    let mut group = c.benchmark_group("query_by_metadata");
    group.sample_size(30);
    let dim = 128;
    let count = 10_000;
    let nodes = generate_nodes(count, dim);

    // JSONL
    {
        let dir = TempDir::new().unwrap();
        let store = create_jsonl_store(&dir);
        populate_backend(&store, &nodes);
        let backend: &dyn StorageBackend = &store;

        group.bench_function("jsonl", |b| {
            b.iter(|| {
                let results = backend
                    .find_by_metadata(
                        "priority",
                        &serde_json::Value::Number(serde_json::Number::from(3)),
                    )
                    .unwrap();
                black_box(results.len())
            });
        });
    }

    // Lance
    #[cfg(feature = "lance-storage")]
    {
        let dir = TempDir::new().unwrap();
        let store = create_lance_store(&dir, dim);
        populate_backend(&store, &nodes);

        group.bench_function("lance", |b| {
            b.iter(|| {
                let results = store
                    .find_by_metadata(
                        "priority",
                        &serde_json::Value::Number(serde_json::Number::from(3)),
                    )
                    .unwrap();
                black_box(results.len())
            });
        });
    }

    group.finish();
}

fn bench_query_composite(c: &mut Criterion) {
    let mut group = c.benchmark_group("query_composite");
    group.sample_size(20);
    let dim = 128;
    let count = 10_000;
    let nodes = generate_nodes(count, dim);

    let filter = QueryFilter {
        tags: Some(vec!["tag-3".to_string()]),
        content_type: Some(ContentType::Text),
        min_importance: Some(0.3),
        metadata: Some(vec![MetadataFilter {
            key: "priority".to_string(),
            value: serde_json::Value::Number(serde_json::Number::from(3)),
        }]),
        limit: 50,
        ..Default::default()
    };

    // JSONL
    {
        let dir = TempDir::new().unwrap();
        let store = create_jsonl_store(&dir);
        populate_backend(&store, &nodes);

        let backend: &dyn StorageBackend = &store;
        group.bench_function("jsonl", |b| {
            b.iter(|| {
                let results = backend.query(&filter).unwrap();
                black_box(results.len())
            });
        });
    }

    // Lance
    #[cfg(feature = "lance-storage")]
    {
        let dir = TempDir::new().unwrap();
        let store = create_lance_store(&dir, dim);
        populate_backend(&store, &nodes);

        group.bench_function("lance", |b| {
            b.iter(|| {
                let results = store.query(&filter).unwrap();
                black_box(results.len())
            });
        });
    }

    group.finish();
}

// ============================================================================
// Benchmark Group 4: Vector Search — Lance ANN vs HNSW
// ============================================================================

fn bench_vector_search(c: &mut Criterion) {
    use aifastdb_core::vector::HNSWIndex;
    use aifastdb_core::types::HNSWConfig;

    let mut group = c.benchmark_group("vector_search");
    group.sample_size(30);
    let dim = 128;

    for count in [1_000, 10_000] {
        let nodes = generate_nodes(count, dim);
        let query = random_embedding(dim);

        // HNSW
        {
            let config = HNSWConfig::new(dim)
                .with_m(16)
                .with_ef_construction(200)
                .with_ef_search(50)
                .with_max_elements(count + 1000);

            let index = HNSWIndex::new(config);
            for node in &nodes {
                if let Some(ref emb) = node.embedding {
                    let _ = index.insert(node.id.clone(), emb.clone());
                }
            }

            group.bench_with_input(
                BenchmarkId::new("hnsw_top10", count),
                &count,
                |b, _| {
                    b.iter(|| {
                        let results = index.search(black_box(&query), 10).unwrap();
                        black_box(results.len())
                    });
                },
            );
        }

        // Lance ANN (if available)
        #[cfg(feature = "lance-storage")]
        {
            let dir = TempDir::new().unwrap();
            let lance_path = dir.path().join("vectors.lance");
            let vec_config = LanceVectorConfig::new(&lance_path, dim);
            let backend = LanceVectorBackend::open(vec_config).unwrap();

            // Insert vectors via VectorBackend trait
            for node in &nodes {
                if let Some(ref emb) = node.embedding {
                    backend.store_vector(&node.id, emb).unwrap();
                }
            }
            backend.flush().unwrap();

            // Create IVF_PQ index if enough vectors
            if count >= 1_000 {
                let index_config = VectorIndexConfig {
                    index_type: VectorIndexType::IvfPq,
                    num_partitions: Some((count as f64).sqrt() as u32),
                    num_sub_vectors: Some(8),
                    metric: VectorDistanceMetric::L2,
                    ..Default::default()
                };
                let _ = backend.create_index(&index_config);
            }

            group.bench_with_input(
                BenchmarkId::new("lance_ann_top10", count),
                &count,
                |b, _| {
                    b.iter(|| {
                        let results = backend.search_nearest(black_box(&query), 10).unwrap_or_default();
                        black_box(results.len())
                    });
                },
            );
        }
    }

    group.finish();
}

// ============================================================================
// Benchmark Group 5: Storage Space & Memory Estimation
// ============================================================================

fn bench_storage_size(c: &mut Criterion) {
    let mut group = c.benchmark_group("storage_size_estimate");
    group.sample_size(10);
    let dim = 128;
    let count = 10_000;
    let nodes = generate_nodes(count, dim);

    // JSONL — measure write time and report file size
    group.bench_function("jsonl_10k_write_and_size", |b| {
        b.iter_with_setup(
            || {
                let dir = TempDir::new().unwrap();
                let store = create_jsonl_store(&dir);
                (dir, store)
            },
            |(dir, store)| {
                for node in &nodes {
                    store.put(node.clone()).unwrap();
                }
                store.flush().unwrap();

                // Measure disk size
                let size = dir_size(dir.path());
                black_box(size)
            },
        );
    });

    // Lance — measure write time and report file size
    #[cfg(feature = "lance-storage")]
    group.bench_function("lance_10k_write_and_size", |b| {
        b.iter_with_setup(
            || {
                let dir = TempDir::new().unwrap();
                let store = create_lance_store(&dir, dim);
                (dir, store)
            },
            |(dir, store)| {
                for node in &nodes {
                    store.put(node.clone()).unwrap();
                }
                store.flush().unwrap();

                // Measure disk size
                let size = dir_size(dir.path());
                black_box(size)
            },
        );
    });

    group.finish();
}

/// Recursively compute the total file size under a directory.
fn dir_size(path: &std::path::Path) -> u64 {
    let mut total = 0u64;
    if path.is_dir() {
        if let Ok(entries) = std::fs::read_dir(path) {
            for entry in entries.flatten() {
                let path = entry.path();
                if path.is_dir() {
                    total += dir_size(&path);
                } else if let Ok(meta) = entry.metadata() {
                    total += meta.len();
                }
            }
        }
    }
    total
}

// ============================================================================
// Benchmark Group 6: Version Control Overhead (Lance-only)
// ============================================================================

#[cfg(feature = "lance-storage")]
fn bench_version_overhead(c: &mut Criterion) {
    let mut group = c.benchmark_group("lance_version_overhead");
    group.sample_size(10);
    let dim = 128;

    // Write 1K nodes, then 10 incremental appends of 100 nodes each
    let base_nodes = generate_nodes(1_000, dim);
    let incremental: Vec<Vec<MemoryNode>> = (0..10)
        .map(|batch_idx| {
            (0..100)
                .map(|i| generate_node(&format!("inc-{}-{}", batch_idx, i), dim))
                .collect()
        })
        .collect();

    group.bench_function("10_incremental_appends", |b| {
        b.iter_with_setup(
            || {
                let dir = TempDir::new().unwrap();
                let store = create_lance_store(&dir, dim);
                populate_backend(&store, &base_nodes);
                (dir, store)
            },
            |(_dir, store)| {
                for batch in &incremental {
                    for node in batch {
                        store.put(node.clone()).unwrap();
                    }
                    store.flush().unwrap();
                }
                black_box(store.len().unwrap())
            },
        );
    });

    // Read after multiple versions
    group.bench_function("read_after_versions", |b| {
        let dir = TempDir::new().unwrap();
        let store = create_lance_store(&dir, dim);
        populate_backend(&store, &base_nodes);
        for batch in &incremental {
            for node in batch {
                store.put(node.clone()).unwrap();
            }
            store.flush().unwrap();
        }

        b.iter(|| {
            let all = store.all().unwrap();
            black_box(all.len())
        });
    });

    group.finish();
}

// Non-lance stub so criterion_group still compiles
#[cfg(not(feature = "lance-storage"))]
fn bench_version_overhead(_c: &mut Criterion) {}

// ============================================================================
// Criterion groups
// ============================================================================

criterion_group!(
    write_benches,
    bench_doc_write_single,
    bench_doc_write_batch,
);

criterion_group!(
    read_benches,
    bench_doc_read_random,
    bench_doc_scan_all,
);

criterion_group!(
    query_benches,
    bench_query_by_tag,
    bench_query_by_metadata,
    bench_query_composite,
);

criterion_group!(
    vector_benches,
    bench_vector_search,
);

criterion_group!(
    storage_benches,
    bench_storage_size,
    bench_version_overhead,
);

criterion_main!(
    write_benches,
    read_benches,
    query_benches,
    vector_benches,
    storage_benches,
);
