//! Standalone benchmark binary that outputs JSON results to stdout.
//!
//! This avoids parsing criterion's internal format and gives a clean JSON
//! that the comparison script can consume.
//!
//! Usage:
//!   cargo bench --package leann-core --bench bench_json_output -- [--skip-large]

use ndarray::Array2;
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use std::collections::BTreeMap;
use std::hint::black_box;
use std::time::Instant;

use leann_core::hnsw::build::{build_hnsw, build_hnsw_with_pool};
use leann_core::hnsw::graph::HnswConfig;
use leann_core::hnsw::io::{read_hnsw_index, write_hnsw_standard};
use leann_core::hnsw::search::{
    SearchBuffers, SearchParams, search_hnsw, search_hnsw_buf, search_hnsw_recompute_buf,
};
use leann_core::hnsw::simd::{inner_product_distance, l2_distance, l2_distance_batch_4};

fn gen_vectors(rng: &mut StdRng, n: usize, d: usize) -> Array2<f32> {
    let data: Vec<f32> = (0..n * d).map(|_| rng.random::<f32>()).collect();
    Array2::from_shape_vec((n, d), data).unwrap()
}

fn gen_query(rng: &mut StdRng, d: usize) -> Vec<f32> {
    (0..d).map(|_| rng.random::<f32>()).collect()
}

/// Run a benchmark: warmup iterations, then `repeats` timed iterations.
/// Returns sorted Vec of elapsed seconds for each iteration.
fn bench_fn<F: FnMut()>(mut f: F, warmup: usize, repeats: usize) -> Vec<f64> {
    // Warmup
    for _ in 0..warmup {
        f();
    }

    // Timed runs
    let mut times = Vec::with_capacity(repeats);
    for _ in 0..repeats {
        let start = Instant::now();
        f();
        times.push(start.elapsed().as_secs_f64());
    }

    times.sort_by(|a, b| a.partial_cmp(b).unwrap());
    times
}

/// Compute a quantile from a sorted slice (linear interpolation).
fn quantile(sorted: &[f64], q: f64) -> f64 {
    if sorted.len() == 1 {
        return sorted[0];
    }
    let pos = q * (sorted.len() - 1) as f64;
    let lo = pos.floor() as usize;
    let hi = lo + 1;
    let frac = pos - lo as f64;
    if hi >= sorted.len() {
        sorted[lo]
    } else {
        sorted[lo] * (1.0 - frac) + sorted[hi] * frac
    }
}

#[derive(serde::Serialize)]
struct BenchResult {
    // Quantiles in seconds
    p5_s: f64,
    p25_s: f64,
    p50_s: f64,
    p75_s: f64,
    p95_s: f64,
    mean_s: f64,
    n_iters: usize,
    // Convenience: median in other units
    median_s: f64,
    median_ms: f64,
    median_us: f64,
    median_ns: f64,
    mean_ms: f64,
    mean_us: f64,
    mean_ns: f64,
}

impl BenchResult {
    fn from_times(times: &[f64]) -> Self {
        let p5 = quantile(times, 0.05);
        let p25 = quantile(times, 0.25);
        let p50 = quantile(times, 0.50);
        let p75 = quantile(times, 0.75);
        let p95 = quantile(times, 0.95);
        let mean = times.iter().sum::<f64>() / times.len() as f64;
        Self {
            p5_s: p5,
            p25_s: p25,
            p50_s: p50,
            p75_s: p75,
            p95_s: p95,
            mean_s: mean,
            n_iters: times.len(),
            median_s: p50,
            median_ms: p50 * 1e3,
            median_us: p50 * 1e6,
            median_ns: p50 * 1e9,
            mean_ms: mean * 1e3,
            mean_us: mean * 1e6,
            mean_ns: mean * 1e9,
        }
    }

    /// Scale all time values by a constant (e.g. to convert batched distance to per-call).
    fn scaled(mut self, factor: f64) -> Self {
        self.p5_s *= factor;
        self.p25_s *= factor;
        self.p50_s *= factor;
        self.p75_s *= factor;
        self.p95_s *= factor;
        self.mean_s *= factor;
        self.median_s *= factor;
        self.median_ms *= factor;
        self.median_us *= factor;
        self.median_ns *= factor;
        self.mean_ms *= factor;
        self.mean_us *= factor;
        self.mean_ns *= factor;
        self
    }

    fn from_size(size: f64) -> Self {
        Self {
            p5_s: size,
            p25_s: size,
            p50_s: size,
            p75_s: size,
            p95_s: size,
            mean_s: size,
            n_iters: 1,
            median_s: size,
            median_ms: size,
            median_us: size,
            median_ns: size,
            mean_ms: size,
            mean_us: size,
            mean_ns: size,
        }
    }
}

fn main() {
    let skip_large = std::env::args().any(|a| a == "--skip-large");
    let num_threads = std::thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(1);
    let hnsw_seed: Option<u64> = std::env::var("HNSW_SEED").ok().and_then(|s| s.parse().ok());
    let mut results: BTreeMap<String, BenchResult> = BTreeMap::new();

    // Create one thread pool, reused across all parallel builds.
    let pool = rayon::ThreadPoolBuilder::new()
        .num_threads(num_threads)
        .build()
        .unwrap();

    eprintln!("Using {num_threads} threads for parallel builds");
    if let Some(seed) = hnsw_seed {
        eprintln!("HNSW_SEED={seed}");
    }

    eprintln!("=== LEANN Rust HNSW Benchmarks (JSON output) ===");

    // ── Distance computation ────────────────────────────────────────
    // Batch multiple calls per timing interval to overcome Instant::now() overhead
    // (~20-50ns on macOS), which would drown out sub-100ns operations.
    eprintln!("\n[1/15] Distance computation...");
    let dist_batch = 1000;
    for dim in [128, 384, 768] {
        let mut rng = StdRng::seed_from_u64(42);
        let a: Vec<f32> = (0..dim).map(|_| rng.random::<f32>()).collect();
        let b: Vec<f32> = (0..dim).map(|_| rng.random::<f32>()).collect();

        let times = bench_fn(
            || {
                for _ in 0..dist_batch {
                    black_box(l2_distance(black_box(&a), black_box(&b)));
                }
            },
            200,
            2_000,
        );
        results.insert(
            format!("distance/l2/{dim}"),
            BenchResult::from_times(&times).scaled(1.0 / dist_batch as f64),
        );

        let times = bench_fn(
            || {
                for _ in 0..dist_batch {
                    black_box(inner_product_distance(black_box(&a), black_box(&b)));
                }
            },
            200,
            2_000,
        );
        results.insert(
            format!("distance/ip/{dim}"),
            BenchResult::from_times(&times).scaled(1.0 / dist_batch as f64),
        );
    }

    // ── HNSW build ──────────────────────────────────────────────────
    eprintln!("[2/15] HNSW build...");
    let config = HnswConfig {
        m: 32,
        ef_construction: 200,
        ef_search: 64,
        distance_metric: leann_core::index::DistanceMetric::L2,
        is_compact: false,
        is_recompute: false,
        seed: hnsw_seed,
    };

    let build_sizes: Vec<usize> = if skip_large {
        vec![100, 1_000, 10_000]
    } else {
        vec![100, 1_000, 10_000, 50_000]
    };

    for &n in &build_sizes {
        eprintln!("  Building {n} vectors...");
        let mut rng = StdRng::seed_from_u64(42);
        let data = gen_vectors(&mut rng, n, 384);

        let repeats = if n >= 50_000 {
            5
        } else if n >= 10_000 {
            10
        } else {
            20
        };
        let times = bench_fn(
            || {
                let _ = build_hnsw_with_pool(&data, &config, &pool, None).unwrap();
            },
            1,
            repeats,
        );
        results.insert(format!("hnsw_build/{n}"), BenchResult::from_times(&times));
    }

    // ── HNSW search (stored vectors) ────────────────────────────────
    eprintln!("[3/15] HNSW search (stored vectors)...");
    {
        let n = 10_000;
        let d = 384;
        let top_k = 10;

        let mut rng = StdRng::seed_from_u64(42);
        let data = gen_vectors(&mut rng, n, d);
        let graph = build_hnsw(&data, &config, None).unwrap();
        let flat_vectors: Vec<f32> = data.iter().copied().collect();
        let query = gen_query(&mut rng, d);

        // Pre-allocate buffers once, reuse across all search calls
        let mut buffers = SearchBuffers::new(graph.ntotal);

        for ef in [16, 32, 64, 128, 256] {
            let params = SearchParams {
                ef_search: ef,
                ..Default::default()
            };
            let times = bench_fn(
                || {
                    let _ = search_hnsw_buf(
                        &graph,
                        &query,
                        top_k,
                        &flat_vectors,
                        &params,
                        &mut buffers,
                    );
                },
                200,
                5000,
            );
            results.insert(
                format!("hnsw_search/ef{ef}"),
                BenchResult::from_times(&times),
            );
        }
    }

    // ── HNSW search (recompute) ─────────────────────────────────────
    eprintln!("[4/15] HNSW search (recompute with in-memory callback)...");
    {
        let n = 10_000;
        let d = 384;
        let top_k = 10;

        let mut rng = StdRng::seed_from_u64(42);
        let data = gen_vectors(&mut rng, n, d);

        let recompute_config = HnswConfig {
            is_recompute: true,
            ..config.clone()
        };
        let graph = build_hnsw(&data, &recompute_config, None).unwrap();
        let flat_vectors: Vec<f32> = data.iter().copied().collect();
        let query = gen_query(&mut rng, d);

        // Pre-allocate buffers once, reuse across all recompute search calls
        let mut recompute_buffers = SearchBuffers::new(graph.ntotal);

        for ef in [16, 32, 64, 128, 256] {
            let params = SearchParams {
                ef_search: ef,
                recompute_embeddings: true,
                ..Default::default()
            };
            let flat_ref = &flat_vectors;
            let times = bench_fn(
                || {
                    let _ = search_hnsw_recompute_buf(
                        &graph,
                        &query,
                        top_k,
                        &params,
                        &mut recompute_buffers,
                        |node_ids, q, out| {
                            let n = node_ids.len();
                            let mut i = 0;
                            while i + 4 <= n {
                                let dists = l2_distance_batch_4(
                                    q,
                                    &flat_ref[node_ids[i] * d..(node_ids[i] + 1) * d],
                                    &flat_ref[node_ids[i + 1] * d..(node_ids[i + 1] + 1) * d],
                                    &flat_ref[node_ids[i + 2] * d..(node_ids[i + 2] + 1) * d],
                                    &flat_ref[node_ids[i + 3] * d..(node_ids[i + 3] + 1) * d],
                                );
                                out[i..i + 4].copy_from_slice(&dists);
                                i += 4;
                            }
                            while i < n {
                                out[i] = l2_distance(
                                    &flat_ref[node_ids[i] * d..(node_ids[i] + 1) * d],
                                    q,
                                );
                                i += 1;
                            }
                        },
                    );
                },
                200,
                5000,
            );
            results.insert(
                format!("hnsw_search_recompute/ef{ef}"),
                BenchResult::from_times(&times),
            );
        }
    }

    // ── Full pipeline (build + write + read + search) ───────────────
    eprintln!("[5/15] Full pipeline...");
    {
        let d = 384;
        let top_k = 10;

        for n in [100, 1_000, 10_000] {
            eprintln!("  Pipeline {n} vectors...");
            let mut rng = StdRng::seed_from_u64(42);
            let data = gen_vectors(&mut rng, n, d);
            let query = gen_query(&mut rng, d);

            let repeats = if n >= 10_000 { 5 } else { 10 };
            let times = bench_fn(
                || {
                    let graph = build_hnsw_with_pool(&data, &config, &pool, None).unwrap();
                    let mut buf = Vec::new();
                    write_hnsw_standard(&mut buf, &graph).unwrap();
                    let mut cursor = std::io::Cursor::new(&buf);
                    let loaded = read_hnsw_index(&mut cursor).unwrap();
                    let flat_vectors: Vec<f32> = data.iter().copied().collect();
                    let params = SearchParams {
                        ef_search: 64,
                        ..Default::default()
                    };
                    let _ = search_hnsw(&loaded, &query, top_k, &flat_vectors, &params);
                },
                1,
                repeats,
            );
            results.insert(
                format!("full_pipeline/{n}"),
                BenchResult::from_times(&times),
            );

            // Also measure index size
            let graph = build_hnsw_with_pool(&data, &config, &pool, None).unwrap();
            let mut buf = Vec::new();
            write_hnsw_standard(&mut buf, &graph).unwrap();
            results.insert(
                format!("index_size_bytes/{n}"),
                BenchResult::from_size(buf.len() as f64),
            );
        }
    }

    // ── Passage lookup ────────────────────────────────────────────────
    eprintln!("[6/15] Passage lookup...");
    {
        use leann_core::index::PassageSource;
        use leann_core::passages::{Passage, PassageManager, write_id_map, write_passages};

        for n in [1_000, 10_000] {
            eprintln!("  Passage lookup n={n}...");
            let dir = tempfile::tempdir().unwrap();
            let passages_path = dir.path().join("test.passages.jsonl");
            let offset_path = dir.path().join("test.passages.idx");
            let id_map_path = dir.path().join("test.ids.txt");

            // Create passages with realistic-ish text
            let passages: Vec<Passage> = (0..n)
                .map(|i| {
                    let mut metadata = std::collections::HashMap::new();
                    metadata.insert("doc_num".to_string(), serde_json::json!(i));
                    metadata.insert(
                        "topic".to_string(),
                        serde_json::json!(format!("topic_{}", i % 10)),
                    );
                    Passage {
                        id: i.to_string(),
                        text: format!(
                            "This is document number {} about topic {}. It contains some text \
                             that is representative of a typical passage in a RAG system.",
                            i,
                            i % 10
                        ),
                        metadata,
                    }
                })
                .collect();

            let ids: Vec<String> = passages.iter().map(|p| p.id.clone()).collect();
            write_passages(&passages, &passages_path, &offset_path).unwrap();
            write_id_map(&ids, &id_map_path).unwrap();

            let sources = vec![PassageSource {
                source_type: "jsonl".to_string(),
                path: passages_path.to_string_lossy().to_string(),
                index_path: offset_path.to_string_lossy().to_string(),
                path_relative: None,
                index_path_relative: None,
            }];

            let manager = PassageManager::load(&sources, None).unwrap();

            // Generate random lookup indices
            let mut rng = StdRng::seed_from_u64(42);
            let lookup_indices: Vec<usize> = (0..1000)
                .map(|_| (rng.random::<f32>() * n as f32) as usize % n)
                .collect();

            let times = bench_fn(
                || {
                    for &idx in &lookup_indices {
                        black_box(manager.get_passage_by_index(idx).unwrap());
                    }
                },
                5,
                100,
            );
            results.insert(
                format!("passage_lookup/{n}"),
                BenchResult::from_times(&times).scaled(1.0 / lookup_indices.len() as f64),
            );
        }
    }

    // ── Text chunking ──────────────────────────────────────────────
    eprintln!("[7/15] Text chunking...");
    {
        use leann_core::chunking::chunk_text;

        let base_sentences = [
            "The quick brown fox jumps over the lazy dog.",
            "Machine learning algorithms can identify patterns in large datasets.",
            "Vector databases enable similarity search across high-dimensional spaces.",
            "Rust provides memory safety without garbage collection.",
            "Neural networks are inspired by the structure of the human brain.",
            "Cloud computing offers scalable infrastructure for modern applications.",
            "Natural language processing has advanced significantly in recent years.",
            "Graph databases model relationships between entities using nodes and edges.",
            "Functional programming emphasizes immutability and pure functions.",
            "The human genome contains approximately three billion base pairs.",
            "Database systems store and retrieve data efficiently using indexing.",
            "Python is a versatile language used for web development and data science.",
            "JavaScript runs in web browsers and is essential for front-end development.",
            "Deep learning has revolutionized computer vision and speech recognition.",
            "Information retrieval systems rank documents by relevance to a query.",
            "Distributed systems coordinate multiple computers to achieve a common goal.",
            "Cryptographic hash functions map arbitrary data to fixed-size outputs.",
            "Operating systems manage hardware resources and provide services to applications.",
            "Compilers translate high-level programming languages into machine code.",
            "Parallel computing divides a problem into subproblems solved simultaneously.",
        ];

        for (label, target_bytes) in [("10KB", 10_000), ("100KB", 100_000), ("1MB", 1_000_000)] {
            eprintln!("  text_chunking/{label}...");
            let mut text = String::with_capacity(target_bytes + 200);
            let mut i = 0;
            while text.len() < target_bytes {
                text.push_str(base_sentences[i % base_sentences.len()]);
                text.push(' ');
                i += 1;
            }

            let times = bench_fn(
                || {
                    black_box(chunk_text(black_box(&text), 512, 50));
                },
                3,
                50,
            );
            results.insert(
                format!("text_chunking/{label}"),
                BenchResult::from_times(&times),
            );
        }
    }

    // ── Sentence splitting ───────────────────────────────────────────
    eprintln!("[8/15] Sentence splitting...");
    {
        use leann_core::chunking::sentence::split_sentences;

        let base_sentences = [
            "The quick brown fox jumps over the lazy dog.",
            "Machine learning algorithms can identify patterns in large datasets.",
            "Vector databases enable similarity search across high-dimensional spaces.",
            "Rust provides memory safety without garbage collection.",
            "Neural networks are inspired by the structure of the human brain.",
            "Cloud computing offers scalable infrastructure for modern applications.",
            "Natural language processing has advanced significantly in recent years.",
            "Graph databases model relationships between entities using nodes and edges.",
            "Functional programming emphasizes immutability and pure functions.",
            "The human genome contains approximately three billion base pairs.",
            "Database systems store and retrieve data efficiently using indexing.",
            "Python is a versatile language used for web development and data science.",
            "JavaScript runs in web browsers and is essential for front-end development.",
            "Deep learning has revolutionized computer vision and speech recognition.",
            "Information retrieval systems rank documents by relevance to a query.",
            "Distributed systems coordinate multiple computers to achieve a common goal.",
            "Cryptographic hash functions map arbitrary data to fixed-size outputs.",
            "Operating systems manage hardware resources and provide services to applications.",
            "Compilers translate high-level programming languages into machine code.",
            "Parallel computing divides a problem into subproblems solved simultaneously.",
        ];

        for (label, target_bytes) in [("10KB", 10_000), ("100KB", 100_000), ("1MB", 1_000_000)] {
            eprintln!("  sentence_split/{label}...");
            let mut text = String::with_capacity(target_bytes + 200);
            let mut i = 0;
            while text.len() < target_bytes {
                text.push_str(base_sentences[i % base_sentences.len()]);
                text.push(' ');
                i += 1;
            }

            let times = bench_fn(
                || {
                    black_box(split_sentences(black_box(&text)));
                },
                3,
                50,
            );
            results.insert(
                format!("sentence_split/{label}"),
                BenchResult::from_times(&times),
            );
        }
    }

    // ── BM25 index build ─────────────────────────────────────────────
    eprintln!("[9/15] BM25 fit...");
    {
        use leann_core::BM25Scorer;

        let topics = [
            "machine learning",
            "database systems",
            "web development",
            "cloud computing",
            "neural networks",
            "programming languages",
            "operating systems",
            "computer vision",
            "natural language",
            "distributed systems",
        ];

        for n in [1_000, 10_000] {
            eprintln!("  bm25_fit/{n}...");
            let docs: Vec<(String, String)> = (0..n)
                .map(|i| {
                    (
                        i.to_string(),
                        format!(
                            "Document {} about {}. The quick brown fox jumps over the lazy dog. \
                             This document discusses various aspects of {} including theory and practice.",
                            i, topics[i % topics.len()], topics[i % topics.len()]
                        ),
                    )
                })
                .collect();

            let repeats = if n >= 10_000 { 10 } else { 30 };
            let times = bench_fn(
                || {
                    let mut scorer = BM25Scorer::default();
                    scorer.fit(black_box(&docs));
                    black_box(&scorer);
                },
                2,
                repeats,
            );
            results.insert(format!("bm25_fit/{n}"), BenchResult::from_times(&times));
        }
    }

    // ── BM25 search ──────────────────────────────────────────────────
    eprintln!("[10/15] BM25 search...");
    {
        use leann_core::BM25Scorer;

        let topics = [
            "machine learning",
            "database systems",
            "web development",
            "cloud computing",
            "neural networks",
            "programming languages",
            "operating systems",
            "computer vision",
            "natural language",
            "distributed systems",
        ];

        for n in [1_000, 10_000] {
            eprintln!("  bm25_search/{n}...");
            let docs: Vec<(String, String)> = (0..n)
                .map(|i| {
                    (
                        i.to_string(),
                        format!(
                            "Document {} about {}. The quick brown fox jumps over the lazy dog. \
                             This document discusses various aspects of {} including theory and practice.",
                            i, topics[i % topics.len()], topics[i % topics.len()]
                        ),
                    )
                })
                .collect();

            let mut scorer = BM25Scorer::default();
            scorer.fit(&docs);

            let times = bench_fn(
                || {
                    black_box(scorer.search(black_box("machine learning neural networks"), 10));
                },
                50,
                500,
            );
            results.insert(format!("bm25_search/{n}"), BenchResult::from_times(&times));
        }
    }

    // ── Metadata filtering ───────────────────────────────────────────
    eprintln!("[11/15] Metadata filtering...");
    {
        use leann_core::{FilterSpec, MetadataFilterEngine, MetadataFilters};

        let genres = ["fiction", "science", "history", "biography", "drama"];
        let topics = [
            "topic_0", "topic_1", "topic_2", "topic_3", "topic_4", "topic_5", "topic_6", "topic_7",
            "topic_8", "topic_9",
        ];

        for n in [1_000, 10_000] {
            eprintln!("  metadata_filter/{n}...");
            let result_set: Vec<std::collections::HashMap<String, serde_json::Value>> = (0..n)
                .map(|i| {
                    let mut m = std::collections::HashMap::new();
                    m.insert("id".to_string(), serde_json::json!(format!("doc{i}")));
                    m.insert(
                        "score".to_string(),
                        serde_json::json!(0.95 - (i as f64 * 0.0001)),
                    );
                    m.insert(
                        "text".to_string(),
                        serde_json::json!(format!("Text for document {i}")),
                    );
                    m.insert(
                        "metadata".to_string(),
                        serde_json::json!({
                            "chapter": (i % 20) + 1,
                            "genre": genres[i % genres.len()],
                            "topic": topics[i % topics.len()],
                            "word_count": 500 + (i * 37) % 5000,
                            "is_published": i % 3 != 0,
                        }),
                    );
                    m
                })
                .collect();

            // Compound filter: chapter <= 5 AND genre in ["fiction", "science"]
            let mut filters = MetadataFilters::new();
            let mut chapter_spec = FilterSpec::new();
            chapter_spec.insert("<=".to_string(), serde_json::json!(5));
            filters.insert("chapter".to_string(), chapter_spec);
            let mut genre_spec = FilterSpec::new();
            genre_spec.insert("in".to_string(), serde_json::json!(["fiction", "science"]));
            filters.insert("genre".to_string(), genre_spec);

            let engine = MetadataFilterEngine::new();

            let times = bench_fn(
                || {
                    black_box(engine.apply_filters(black_box(&result_set), black_box(&filters)));
                },
                10,
                200,
            );
            results.insert(
                format!("metadata_filter/{n}"),
                BenchResult::from_times(&times),
            );
        }
    }

    // ── Index I/O (write) ────────────────────────────────────────────
    eprintln!("[12/15] Index I/O (write)...");
    {
        for n in [1_000, 10_000] {
            eprintln!("  index_io_write/{n}...");
            let mut rng = StdRng::seed_from_u64(42);
            let data = gen_vectors(&mut rng, n, 384);
            let graph = build_hnsw_with_pool(&data, &config, &pool, None).unwrap();

            let times = bench_fn(
                || {
                    let mut buf = Vec::new();
                    write_hnsw_standard(&mut buf, &graph).unwrap();
                    black_box(buf);
                },
                5,
                50,
            );
            results.insert(
                format!("index_io_write/{n}"),
                BenchResult::from_times(&times),
            );
        }
    }

    // ── Index I/O (read) ─────────────────────────────────────────────
    eprintln!("[13/15] Index I/O (read)...");
    {
        for n in [1_000, 10_000] {
            eprintln!("  index_io_read/{n}...");
            let mut rng = StdRng::seed_from_u64(42);
            let data = gen_vectors(&mut rng, n, 384);
            let graph = build_hnsw_with_pool(&data, &config, &pool, None).unwrap();

            let mut buf = Vec::new();
            write_hnsw_standard(&mut buf, &graph).unwrap();

            let times = bench_fn(
                || {
                    let mut cursor = std::io::Cursor::new(&buf);
                    let loaded = read_hnsw_index(&mut cursor).unwrap();
                    black_box(loaded);
                },
                5,
                50,
            );
            results.insert(
                format!("index_io_read/{n}"),
                BenchResult::from_times(&times),
            );
        }
    }

    // ── Passage file write ───────────────────────────────────────────
    eprintln!("[14/15] Passage write...");
    {
        use leann_core::passages::{Passage, write_passages};

        for n in [1_000usize, 10_000] {
            eprintln!("  passage_write/{n}...");
            let passages: Vec<Passage> = (0..n)
                .map(|i| {
                    let mut metadata = std::collections::HashMap::new();
                    metadata.insert("doc_num".to_string(), serde_json::json!(i));
                    metadata.insert(
                        "topic".to_string(),
                        serde_json::json!(format!("topic_{}", i % 10)),
                    );
                    Passage {
                        id: i.to_string(),
                        text: format!(
                            "This is document number {} about topic {}. It contains some text \
                             that is representative of a typical passage in a RAG system.",
                            i,
                            i % 10
                        ),
                        metadata,
                    }
                })
                .collect();

            let dir = tempfile::tempdir().unwrap();

            let times = bench_fn(
                || {
                    let p_path = dir.path().join("bench.passages.jsonl");
                    let o_path = dir.path().join("bench.passages.idx");
                    write_passages(&passages, &p_path, &o_path).unwrap();
                },
                2,
                30,
            );
            results.insert(
                format!("passage_write/{n}"),
                BenchResult::from_times(&times),
            );
        }
    }

    // ── CLI startup time ─────────────────────────────────────────────
    eprintln!("[15/15] CLI startup time...");
    {
        // Find the Rust CLI binary (release build).
        // cargo bench runs with cwd = crate dir, so walk up to find the workspace target/.
        let rust_bin = std::env::current_dir()
            .unwrap()
            .ancestors()
            .map(|p| p.join("target/release/leann"))
            .find(|p| p.exists())
            .unwrap_or_else(|| {
                std::env::current_dir()
                    .unwrap()
                    .join("target/release/leann")
            });
        if rust_bin.exists() {
            eprintln!("  Measuring Rust CLI startup ({})...", rust_bin.display());
            let mut times = Vec::with_capacity(20);
            // Warmup
            for _ in 0..3 {
                let _ = std::process::Command::new(&rust_bin)
                    .arg("--help")
                    .stdout(std::process::Stdio::null())
                    .stderr(std::process::Stdio::null())
                    .status();
            }
            for _ in 0..20 {
                let start = Instant::now();
                let _ = std::process::Command::new(&rust_bin)
                    .arg("--help")
                    .stdout(std::process::Stdio::null())
                    .stderr(std::process::Stdio::null())
                    .status();
                times.push(start.elapsed().as_secs_f64());
            }
            times.sort_by(|a, b| a.partial_cmp(b).unwrap());
            results.insert("cli_startup".to_string(), BenchResult::from_times(&times));
        } else {
            eprintln!(
                "  SKIP: Rust CLI binary not found at {}. Run `cargo build --release -p leann-cli` first.",
                rust_bin.display()
            );
        }
    }

    eprintln!("\nDone! Writing JSON to stdout.");
    println!("{}", serde_json::to_string_pretty(&results).unwrap());
}
