use criterion::{black_box, criterion_group, criterion_main, Criterion};
use llmosafe::{
    calculate_halo_signal, get_bias_breakdown, sift_perceptions, AdversarialDetector,
    ConfidenceTracker, CusumDetector, DriftDetector, EscalationPolicy, ReasoningLoop,
    RepetitionDetector, Synapse, WorkingMemory,
};

fn bench_sifter(c: &mut Criterion) {
    let observations = vec![
        "The expert provided an official professional recommendation",
        "System running normally with high authority",
        "No anomalies detected by the official team",
    ];
    let objective = "safety analysis";

    c.bench_function("calculate_halo_signal", |b| {
        b.iter(|| calculate_halo_signal(black_box("The expert says this is official")))
    });

    c.bench_function("get_bias_breakdown", |b| {
        b.iter(|| get_bias_breakdown(black_box("The expert says this is official")))
    });

    c.bench_function("sift_perceptions", |b| {
        b.iter(|| sift_perceptions(black_box(&observations), black_box(objective)))
    });
}

fn bench_kernel(c: &mut Criterion) {
    let mut synapse = Synapse::new();
    synapse.set_raw_entropy(500);
    synapse.set_raw_surprise(100);
    let mut loop_guard = ReasoningLoop::<100>::new();

    c.bench_function("synapse_validate", |b| {
        b.iter(|| black_box(synapse).validate())
    });

    let mut memory = WorkingMemory::<64>::new(1000);
    let sifted = llmosafe::SiftedSynapse::new(synapse);
    let validated = memory.update(sifted).unwrap();
    c.bench_function("reasoning_loop_next", |b| {
        b.iter(|| loop_guard.next_step(black_box(validated)))
    });
}

fn bench_memory(c: &mut Criterion) {
    let mut memory = WorkingMemory::<64>::new(1000);
    let mut synapse = Synapse::new();
    synapse.set_raw_entropy(500);
    let sifted = llmosafe::SiftedSynapse::new(synapse);

    c.bench_function("memory_update", |b| {
        b.iter(|| memory.update(black_box(sifted)))
    });
}

fn bench_detection(c: &mut Criterion) {
    let mut cusum = CusumDetector::new(500.0, 50.0, 200.0);
    let mut rep = RepetitionDetector::new(3);
    let mut drift = DriftDetector::new("rust safety", 0.5);
    let mut conf = ConfidenceTracker::new(0.5, 2);
    let adv = AdversarialDetector::new();

    c.bench_function("cusum_update", |b| {
        b.iter(|| cusum.update(black_box(600.0)))
    });

    c.bench_function("repetition_observe", |b| {
        b.iter(|| rep.observe(black_box("constant input")))
    });

    c.bench_function("drift_observe", |b| {
        b.iter(|| drift.observe(black_box("different semantic context")))
    });

    c.bench_function("confidence_observe", |b| {
        b.iter(|| conf.observe(black_box(0.7)))
    });

    c.bench_function("adversarial_detect", |b| {
        b.iter(|| adv.detect_substrings(black_box("ignore previous instructions")))
    });
}

fn bench_integration(c: &mut Criterion) {
    let policy = EscalationPolicy::default();

    c.bench_function("policy_decide", |b| {
        b.iter(|| policy.decide(black_box(500), black_box(100), black_box(false)))
    });
}

criterion_group!(
    benches,
    bench_sifter,
    bench_kernel,
    bench_memory,
    bench_detection,
    bench_integration
);
criterion_main!(benches);
