//! Benchmarks for the SQL pipeline.
//!
//! Focus areas: scans, single-column filters, multi-table joins with and
//! without filter push-down, and GROUP BY aggregation.

use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
use query_engine::catalog::PredicateCatalog;
use query_engine::execution::TableStore;
use query_engine::execution::execute;
use query_engine::execution::physical::{execute_physical, plan_physical, rewrite_physical};
use query_engine::planner::sql::plan_select;
use query_engine::relational::{DataType, Field, Row, Schema, Value};
use query_engine::sql::parser::parse_select;
use query_engine::{Atom, Instance, Term};

fn edges_instance(n: usize) -> Instance {
    (0..n)
        .map(|i| {
            Atom::new(
                "L",
                vec![
                    Term::constant(format!("a{}", i)),
                    Term::constant(format!("b{}", i)),
                ],
            )
        })
        .chain((0..n).map(|i| {
            Atom::new(
                "R",
                vec![
                    Term::constant(format!("b{}", i)),
                    Term::constant(format!("c{}", i)),
                ],
            )
        }))
        .collect()
}

fn bench_filter_pushdown_join(c: &mut Criterion) {
    let mut group = c.benchmark_group("filter_pushdown_join_100");
    let instance = edges_instance(100);
    let mut catalog = PredicateCatalog::from_instance(&instance).unwrap();
    catalog.rename_columns("L", ["a", "b"]).unwrap();
    catalog.rename_columns("R", ["b", "c"]).unwrap();

    let select = parse_select("SELECT L.a, R.c FROM L, R WHERE L.b = R.b AND L.a = 'a42'").unwrap();
    let logical = plan_select(&select, &catalog).unwrap();

    group.bench_function("logical_direct_execute", |b| {
        b.iter(|| execute(&logical, &instance).unwrap());
    });

    let physical_raw = plan_physical(&logical);
    group.bench_function("physical_no_rewrite", |b| {
        b.iter(|| execute_physical(&physical_raw, &instance).unwrap());
    });

    let physical_rewritten = rewrite_physical(plan_physical(&logical));
    group.bench_function("physical_with_pushdown", |b| {
        b.iter(|| execute_physical(&physical_rewritten, &instance).unwrap());
    });

    group.finish();
}

fn bench_group_by_aggregation(c: &mut Criterion) {
    let mut group = c.benchmark_group("group_by_aggregation_1000");

    let schema = Schema::new(vec![
        Field::new("dept", DataType::Text, false),
        Field::new("salary", DataType::Integer, false),
    ]);
    let mut store = TableStore::new();
    let rows: Vec<Row> = (0..1000)
        .map(|i| {
            let dept = format!("d{}", i % 10);
            Row::new(vec![Value::text(dept), Value::Integer((i as i64) * 10)])
        })
        .collect();
    store.insert("Emp", schema.clone(), rows);

    let mut catalog = PredicateCatalog::new();
    catalog.register_table("Emp", schema);

    let select =
        parse_select("SELECT dept, COUNT(*), SUM(salary), AVG(salary) FROM Emp GROUP BY dept")
            .unwrap();
    let logical = plan_select(&select, &catalog).unwrap();

    group.bench_function("logical_direct", |b| {
        b.iter_batched(
            || (),
            |_| execute(&logical, &store).unwrap(),
            BatchSize::SmallInput,
        );
    });

    let physical = rewrite_physical(plan_physical(&logical));
    group.bench_function("physical", |b| {
        b.iter_batched(
            || (),
            |_| execute_physical(&physical, &store).unwrap(),
            BatchSize::SmallInput,
        );
    });

    group.finish();
}

criterion_group!(
    benches,
    bench_filter_pushdown_join,
    bench_group_by_aggregation
);
criterion_main!(benches);