use criterion::{Criterion, black_box, criterion_group, criterion_main};

use keydr::engine::key_stats::KeyStatsStore;
use keydr::engine::ngram_stats::{BigramKey, BigramStatsStore, extract_ngram_events};
use keydr::session::result::KeyTime;

fn make_keystrokes(count: usize) -> Vec<KeyTime> {
    let chars = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'];
    (0..count)
        .map(|i| KeyTime {
            key: chars[i % chars.len()],
            time_ms: 200.0 + (i % 50) as f64,
            correct: i % 7 != 0, // ~14% error rate
        })
        .collect()
}

fn bench_extraction(c: &mut Criterion) {
    let keystrokes = make_keystrokes(500);

    c.bench_function("extract_bigrams (500 keystrokes)", |b| {
        b.iter(|| extract_ngram_events(black_box(&keystrokes), 800.0))
    });
}

fn bench_update(c: &mut Criterion) {
    let keystrokes = make_keystrokes(500);
    let bigram_events = extract_ngram_events(&keystrokes, 800.0);

    c.bench_function("bigram_stats update (400 events)", |b| {
        b.iter(|| {
            let mut store = BigramStatsStore::default();
            for ev in bigram_events.iter().take(400) {
                store.update(
                    black_box(ev.key.clone()),
                    black_box(ev.total_time_ms),
                    black_box(ev.correct),
                    black_box(ev.has_hesitation),
                    0,
                );
            }
            store
        })
    });
}

fn bench_focus_selection(c: &mut Criterion) {
    // Use a-z + A-Z + 0-9 = 62 chars for up to 3844 unique bigrams
    let all_chars: Vec<char> = ('a'..='z').chain('A'..='Z').chain('0'..='9').collect();

    let mut bigram_stats = BigramStatsStore::default();
    let mut char_stats = KeyStatsStore::default();

    // Char baselines: 5% error rate, 430ms per char. Expected independent bigram
    // error rate is therefore 1 - 0.95^2 = 0.0975, so a bigram needs an error_rate_ema
    // above ~0.146 to clear the 1.5x error-anomaly ratio threshold, and a
    // filtered_time_ms above ~645 to clear the 50% speed-anomaly threshold.
    for &ch in &all_chars {
        let stat = char_stats.stats.entry(ch).or_default();
        stat.confidence = 0.8;
        stat.filtered_time_ms = 430.0;
        stat.sample_count = 50;
        stat.total_count = 50;
        stat.error_count = 3;
        stat.error_rate_ema = 0.05;
    }

    let mut count: usize = 0;
    for &a in &all_chars {
        for &b in &all_chars {
            if bigram_stats.stats.len() >= 3000 {
                break;
            }
            let key = BigramKey([a, b]);
            let stat = bigram_stats.stats.entry(key).or_default();
            // Spread values so roughly half the entries clear each anomaly threshold,
            // exercising both the error and speed candidate paths.
            stat.error_rate_ema = 0.05 + (count % 20) as f64 * 0.01;
            stat.filtered_time_ms = 500.0 + (count % 40) as f64 * 10.0;
            stat.sample_count = 25 + count % 30;
            stat.error_count = 5 + count % 10;
            // Streak >= 3 confirms an anomaly; mix confirmed and unconfirmed entries.
            stat.error_anomaly_streak = if count % 3 == 0 { 3 } else { 1 };
            stat.speed_anomaly_streak = if count % 4 == 0 { 3 } else { 0 };
            count += 1;
        }
    }
    assert_eq!(bigram_stats.stats.len(), 3000);

    let unlocked: Vec<char> = all_chars;

    c.bench_function("worst_confirmed_anomaly (3K entries)", |b| {
        b.iter(|| {
            bigram_stats.worst_confirmed_anomaly(black_box(&char_stats), black_box(&unlocked))
        })
    });
}

fn bench_history_replay(c: &mut Criterion) {
    // Build 500 drills of ~300 keystrokes each
    let drills: Vec<Vec<KeyTime>> = (0..500).map(|_| make_keystrokes(300)).collect();

    c.bench_function("history replay (500 drills x 300 keystrokes)", |b| {
        b.iter(|| {
            let mut bigram_stats = BigramStatsStore::default();
            let mut key_stats = KeyStatsStore::default();

            for (drill_idx, keystrokes) in drills.iter().enumerate() {
                let bigram_events = extract_ngram_events(keystrokes, 800.0);

                for kt in keystrokes {
                    if kt.correct {
                        let stat = key_stats.stats.entry(kt.key).or_default();
                        stat.total_count += 1;
                    } else {
                        key_stats.update_key_error(kt.key);
                    }
                }

                for ev in &bigram_events {
                    bigram_stats.update(
                        ev.key.clone(),
                        ev.total_time_ms,
                        ev.correct,
                        ev.has_hesitation,
                        drill_idx as u32,
                    );
                }
            }

            (bigram_stats, key_stats)
        })
    });
}

criterion_group!(
    benches,
    bench_extraction,
    bench_update,
    bench_focus_selection,
    bench_history_replay,
);
criterion_main!(benches);