use criterion::{black_box, criterion_group, criterion_main, Criterion};
use klieo_core::ids::RunId;
use klieo_core::memory::Episode;
use klieo_memory_sqlite::{DummyEmbedder, MemorySqlite};
use std::sync::Arc;

// Shared runtime — re-creating the runtime between criterion benches leaves
// deadpool-sqlite without a reactor and panics. Keep one runtime alive for
// every bench in this file.
fn shared_rt() -> &'static tokio::runtime::Runtime {
    use std::sync::OnceLock;
    static RT: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
    RT.get_or_init(|| {
        tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .enable_all()
            .build()
            .expect("rt")
    })
}

fn sample_event() -> Episode {
    Episode::llm_call(1280, 250)
}

fn bench_record_single(c: &mut Criterion) {
    let rt = shared_rt();
    let _enter = rt.enter();
    let mem = rt.block_on(async {
        MemorySqlite::new(":memory:", Arc::new(DummyEmbedder))
            .await
            .expect("open")
    });
    let run = RunId::new();
    c.bench_function("sqlite_episodic_record_single", |b| {
        b.iter(|| {
            rt.block_on(async {
                mem.episodic
                    .record(run, black_box(sample_event()))
                    .await
                    .expect("record");
            });
        });
    });
}

fn bench_record_amortized(c: &mut Criterion) {
    // Detects O(E^2) per-run scaling: cost of recording into a long-lived run.
    let rt = shared_rt();
    let _enter = rt.enter();
    let mut group = c.benchmark_group("sqlite_episodic_record_amortized");
    for &preseed in &[0usize, 100, 1000] {
        let mem = rt.block_on(async {
            MemorySqlite::new(":memory:", Arc::new(DummyEmbedder))
                .await
                .expect("open")
        });
        let run = RunId::new();
        rt.block_on(async {
            for _ in 0..preseed {
                mem.episodic
                    .record(run, sample_event())
                    .await
                    .expect("seed");
            }
        });
        group.bench_function(format!("preseed_{preseed}"), |b| {
            b.iter(|| {
                rt.block_on(async {
                    mem.episodic
                        .record(run, black_box(sample_event()))
                        .await
                        .expect("record");
                });
            });
        });
    }
    group.finish();
}

criterion_group!(benches, bench_record_single, bench_record_amortized);
criterion_main!(benches);
