//! skade throughput benchmark — append / partitioned-append / bulk-ingest /
//! full-scan over an embedded warehouse. Mirrors the write/read shapes the
//! nornir-catalog bench (and nornir's warehouse) drive through skade.
//!
//! Each append is one `fast_append` snapshot == one redb commit + fsync, so
//! these numbers are commit-bound at small batch sizes (that's the real cost of
//! the per-call API); `ingest` amortises the commit across many batches.
//!
//! Run: `cargo bench -p skade`  (or `cargo bench --bench throughput`)

use std::sync::Arc;

use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use skade::arrow_array::{Float64Array, Int64Array, RecordBatch, StringArray};
use skade::arrow_schema::{DataType, Field, Schema};
use tokio::runtime::Runtime;

fn schema() -> Arc<Schema> {
    Arc::new(Schema::new(vec![
        Field::new("repo", DataType::Utf8, false),
        Field::new("id", DataType::Int64, false),
        Field::new("name", DataType::Utf8, false),
        Field::new("score", DataType::Float64, false),
    ]))
}

fn batch(rows: usize) -> RecordBatch {
    RecordBatch::try_new(
        schema(),
        vec![
            Arc::new(StringArray::from(vec!["nornir"; rows])),
            Arc::new(Int64Array::from((0..rows as i64).collect::<Vec<_>>())),
            Arc::new(StringArray::from((0..rows).map(|i| format!("row-{i}")).collect::<Vec<_>>())),
            Arc::new(Float64Array::from((0..rows).map(|i| i as f64).collect::<Vec<_>>())),
        ],
    )
    .unwrap()
}

fn rt() -> Runtime {
    Runtime::new().unwrap()
}

/// Single-batch `append` (one snapshot/commit per call) — commit-bound.
fn bench_append(c: &mut Criterion) {
    let rt = rt();
    let mut g = c.benchmark_group("append");
    g.sample_size(20);
    for &rows in &[1_000usize, 10_000] {
        let b = batch(rows);
        g.throughput(Throughput::Elements(rows as u64));
        g.bench_with_input(BenchmarkId::from_parameter(rows), &rows, |bench, _| {
            let tmp = tempfile::tempdir().unwrap();
            let mut table = rt.block_on(async {
                let wh = skade::open(tmp.path()).await.unwrap();
                wh.create_table("t", schema().as_ref()).await.unwrap()
            });
            bench.iter(|| {
                rt.block_on(async { table.append(std::slice::from_ref(&b)).await.unwrap() });
            });
        });
    }
    g.finish();
}

/// Partitioned (`repo`) single-batch append — adds skade's per-commit
/// identity `partition_key_for` over the unpartitioned path.
fn bench_append_partitioned(c: &mut Criterion) {
    let rt = rt();
    let mut g = c.benchmark_group("append_partitioned_repo");
    g.sample_size(20);
    let rows = 10_000usize;
    let b = batch(rows);
    g.throughput(Throughput::Elements(rows as u64));
    g.bench_function(BenchmarkId::from_parameter(rows), |bench| {
        let tmp = tempfile::tempdir().unwrap();
        let mut table = rt.block_on(async {
            let wh = skade::open(tmp.path()).await.unwrap();
            wh.create_partitioned_table("t", schema().as_ref(), &["repo"]).await.unwrap()
        });
        bench.iter(|| {
            rt.block_on(async { table.append(std::slice::from_ref(&b)).await.unwrap() });
        });
    });
    g.finish();
}

/// Bulk `ingest` — many batches, commit every `batches_per_commit`. Amortises
/// the fsync, so this is the real bulk-load throughput.
fn bench_ingest(c: &mut Criterion) {
    let rt = rt();
    let mut g = c.benchmark_group("ingest");
    g.sample_size(10);
    let (n_batches, rows) = (64usize, 1_000usize);
    g.throughput(Throughput::Elements((n_batches * rows) as u64));
    g.bench_function("64x1000_commit8", |bench| {
        bench.iter_batched(
            || {
                // Fresh warehouse + table per iteration (ingest mutates state).
                let tmp = tempfile::tempdir().unwrap();
                let table = rt.block_on(async {
                    let wh = skade::open(tmp.path()).await.unwrap();
                    wh.create_table("t", schema().as_ref()).await.unwrap()
                });
                let batches: Vec<RecordBatch> = (0..n_batches).map(|_| batch(rows)).collect();
                (tmp, table, batches)
            },
            |(tmp, mut table, batches)| {
                rt.block_on(async { table.ingest(batches, 8).await.unwrap() });
                drop(tmp);
            },
            criterion::BatchSize::SmallInput,
        );
    });
    g.finish();
}

/// Full-snapshot `read` (scan → Arrow) over a pre-populated table.
fn bench_read(c: &mut Criterion) {
    let rt = rt();
    let mut g = c.benchmark_group("read");
    g.sample_size(20);
    let (n_batches, rows) = (100usize, 1_000usize);
    let total = (n_batches * rows) as u64;

    let tmp = tempfile::tempdir().unwrap();
    let table = rt.block_on(async {
        let wh = skade::open(tmp.path()).await.unwrap();
        let mut t = wh.create_table("t", schema().as_ref()).await.unwrap();
        let batches: Vec<RecordBatch> = (0..n_batches).map(|_| batch(rows)).collect();
        t.ingest(batches, 16).await.unwrap();
        t
    });

    g.throughput(Throughput::Elements(total));
    g.bench_function("scan_100k", |bench| {
        bench.iter(|| {
            let got = rt.block_on(async { table.read().await.unwrap() });
            assert_eq!(got.iter().map(|b| b.num_rows()).sum::<usize>(), (n_batches * rows));
        });
    });
    g.finish();
}

/// Parallel-encode scaling — the SAME `ingest_parallel` run under a 1-thread
/// rayon pool vs the all-cores default pool. Isolates the multi-core *encode*
/// win (the gatling payoff): writes/commits are identical, only the Parquet
/// encoding fans out. Encode-heavy on purpose (32 files × 20k rows) so the CPU
/// work dominates the per-commit fsync.
fn bench_ingest_parallel_scaling(c: &mut Criterion) {
    let rt = rt();
    let mut g = c.benchmark_group("ingest_parallel_scaling");
    g.sample_size(10);
    let (n_files, rows) = (32usize, 20_000usize);
    g.throughput(Throughput::Elements((n_files * rows) as u64));

    for (label, threads) in [("1core", 1usize), ("allcores", 0usize)] {
        let pool = if threads == 0 {
            rayon::ThreadPoolBuilder::new().build().unwrap() // default = all cores
        } else {
            rayon::ThreadPoolBuilder::new().num_threads(threads).build().unwrap()
        };
        g.bench_function(label, |bench| {
            bench.iter_batched(
                || {
                    let tmp = tempfile::tempdir().unwrap();
                    let table = rt.block_on(async {
                        let wh = skade::open(tmp.path()).await.unwrap();
                        wh.create_table("t", schema().as_ref()).await.unwrap()
                    });
                    let groups: Vec<Vec<RecordBatch>> =
                        (0..n_files).map(|_| vec![batch(rows)]).collect();
                    (tmp, table, groups)
                },
                |(tmp, mut table, groups)| {
                    // par_iter inside ingest_parallel runs on the installed pool.
                    pool.install(|| {
                        rt.block_on(async { table.ingest_parallel(groups, 8).await.unwrap() });
                    });
                    drop(tmp);
                },
                criterion::BatchSize::SmallInput,
            );
        });
    }
    g.finish();
}

criterion_group!(
    benches,
    bench_append,
    bench_append_partitioned,
    bench_ingest,
    bench_read,
    bench_ingest_parallel_scaling
);
criterion_main!(benches);
