use std::{future::Future, pin::Pin, sync::OnceLock};

use criterion::{black_box, criterion_group, criterion_main, Criterion};
use incria::{
    deps,
    thunk::{Thunk, ThunkMapper},
    Mapper,
};

type Key = (usize, usize);
type Value = usize;

#[derive(Debug, Default)]
struct CellThunk;

impl Thunk<Key, Value> for CellThunk {
    fn compute(&self, key: Key) -> Pin<Box<dyn Future<Output = Value> + Send + '_>> {
        Box::pin(async move {
            if key.0 > 0 {
                let prev1 = *cell_mapping().get(&(key.0 - 1, key.1 - 1)).await;
                let prev2 = *cell_mapping().get(&(key.0 - 1, key.1)).await;
                let prev3 = *cell_mapping().get(&(key.0 - 1, key.1 + 1)).await;
                let val = prev1 + prev2 + prev3;

                if val > 10_000 {
                    1
                } else {
                    val
                }
            } else {
                1
            }
        })
    }
}

fn calc_raw(key: Key) -> Value {
    if key.0 > 0 {
        let prev1 = calc_raw((key.0 - 1, key.1 - 1));
        let prev2 = calc_raw((key.0 - 1, key.1));
        let prev3 = calc_raw((key.0 - 1, key.1 + 1));
        let val = prev1 + prev2 + prev3;

        if val > 10_000 {
            1
        } else {
            val
        }
    } else {
        1
    }
}

type CellMapping = ThunkMapper<Key, Value, CellThunk>;
static CELL_MAPPING: OnceLock<CellMapping> = OnceLock::new();
fn cell_mapping() -> &'static CellMapping {
    CELL_MAPPING.get_or_init(CellMapping::default)
}

fn criterion_benchmark(c: &mut Criterion) {
    c.bench_function("spreadsheet non-incria (n = 5)", |b| {
        b.iter(|| calc_raw(black_box((5, 5))));
    });
    c.bench_function("spreadsheet fresh (n = 5)", |b| {
        b.to_async(tokio::runtime::Runtime::new().unwrap())
            .iter(|| deps::with_node_id(deps::next_node_id(), cell_mapping().get(&(5, 5))));
    });
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);