use criterion::{Criterion, criterion_group, criterion_main};
use mimirs_agent::MimirAgent;
use mimirs_core::{MemoryAgent, MemoryInput, RecallQuery};
use mimirs_store::Urdarbrunnr;
use std::sync::Arc;
use tempfile::tempdir;
use tokio::runtime::Runtime;

fn bench_agent_load(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let dir = tempdir().unwrap();
    let well = Arc::new(Urdarbrunnr::opna(dir.path(), 384).unwrap());
    let agent = MimirAgent::new(well, mimirs_agent::AgentConfig::default());

    c.bench_function("agent_rista_rista_rista", |b| {
        b.to_async(&rt).iter(|| async {
            let input = MemoryInput {
                content: "Production-like load test entry with sufficient entropy for density matrix generation.".to_string(),
                metadata: None,
                verifiability: None,
                scope: None,
                memory_class: None,
                source_session: None,
            };
            agent.store(input).await.unwrap();
        });
    });

    c.bench_function("agent_minni_recall", |b| {
        // Pre-populate with some memories
        rt.block_on(async {
            for i in 0..100 {
                let input = MemoryInput {
                    content: format!("Memory entry index {}", i),
                    metadata: None,
                    verifiability: None,
                    scope: None,
                    memory_class: None,
                    source_session: None,
                };
                agent.store(input).await.unwrap();
            }
        });

        b.to_async(&rt).iter(|| async {
            let query = RecallQuery::new("Memory entry").with_limit(10);
            agent.recall(query).await.unwrap();
        });
    });
}

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