#![allow(clippy::unwrap_used)]
//! Benchmark: MLS group operations at various group sizes.
//!
//! Measures KeyPackage generation, group creation, member addition,
//! message encryption, and message decryption.

use std::sync::Arc;

use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
use quicprochat_core::{GroupMember, IdentityKeypair};

/// Create identities and a group of the given size.
/// Returns (creator, Vec<members>).
fn setup_group(size: usize) -> (GroupMember, Vec<GroupMember>) {
    let creator_id = Arc::new(IdentityKeypair::generate());
    let mut creator = GroupMember::new(creator_id);
    creator.create_group(b"bench-group").unwrap();

    let mut members = Vec::with_capacity(size.saturating_sub(1));
    for _ in 1..size {
        let joiner_id = Arc::new(IdentityKeypair::generate());
        let mut joiner = GroupMember::new(joiner_id);
        let kp = joiner.generate_key_package().unwrap();

        let (_commit, welcome) = creator.add_member(&kp).unwrap();
        joiner.join_group(&welcome).unwrap();
        members.push(joiner);
    }

    (creator, members)
}

fn bench_keygen(c: &mut Criterion) {
    c.bench_function("mls_keygen", |b| {
        b.iter_batched(
            || {
                let id = Arc::new(IdentityKeypair::generate());
                GroupMember::new(id)
            },
            |mut member| {
                member.generate_key_package().unwrap();
            },
            BatchSize::SmallInput,
        );
    });
}

fn bench_group_create(c: &mut Criterion) {
    c.bench_function("mls_group_create", |b| {
        b.iter_batched(
            || {
                let id = Arc::new(IdentityKeypair::generate());
                GroupMember::new(id)
            },
            |mut member| {
                member.create_group(b"bench-group").unwrap();
            },
            BatchSize::SmallInput,
        );
    });
}

fn bench_add_member(c: &mut Criterion) {
    let mut group = c.benchmark_group("mls_add_member");
    group.sample_size(10);
    for size in [2, 10, 50, 100] {
        group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &size| {
            b.iter_batched(
                || {
                    let (creator, members) = setup_group(size);
                    let joiner_id = Arc::new(IdentityKeypair::generate());
                    let mut joiner = GroupMember::new(joiner_id);
                    let kp = joiner.generate_key_package().unwrap();
                    (creator, members, joiner, kp)
                },
                |(mut creator, _members, _joiner, kp)| {
                    creator.add_member(&kp).unwrap();
                },
                BatchSize::SmallInput,
            );
        });
    }
    group.finish();
}

fn bench_epoch_rotation(c: &mut Criterion) {
    let mut group = c.benchmark_group("mls_epoch_rotation");
    group.sample_size(10);
    for size in [2, 10, 50] {
        group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &size| {
            b.iter_batched(
                || {
                    let (mut creator, members) = setup_group(size);
                    // Propose a self-update to simulate epoch rotation
                    let proposal = creator.propose_self_update().unwrap();
                    (creator, members, proposal)
                },
                |(mut creator, _members, _proposal)| {
                    // Commit pending proposals (the self-update) to advance the epoch
                    creator.commit_pending_proposals().unwrap();
                },
                BatchSize::SmallInput,
            );
        });
    }
    group.finish();
}

fn bench_send_message(c: &mut Criterion) {
    let mut group = c.benchmark_group("mls_send_message");
    for size in [2, 10, 50] {
        group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &size| {
            let (mut creator, _members) = setup_group(size);
            let payload = b"hello benchmark message";
            b.iter(|| {
                creator.send_message(payload).unwrap();
            });
        });
    }
    group.finish();
}

fn bench_receive_message(c: &mut Criterion) {
    let mut group = c.benchmark_group("mls_receive_message");
    for size in [2, 10, 50] {
        group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &size| {
            // For receive, we need a fresh ciphertext each iteration since
            // MLS message processing is destructive (epoch state changes).
            // We pre-generate a batch and consume them.
            let (mut creator, mut members) = setup_group(size);
            if members.is_empty() {
                return;
            }
            let payload = b"hello benchmark message";
            b.iter_batched(
                || creator.send_message(payload).unwrap(),
                |ct| {
                    // Receive on the first joiner
                    let _ = members[0].receive_message(&ct);
                },
                BatchSize::SmallInput,
            );
        });
    }
    group.finish();
}

criterion_group!(
    benches,
    bench_keygen,
    bench_group_create,
    bench_add_member,
    bench_epoch_rotation,
    bench_send_message,
    bench_receive_message,
);
criterion_main!(benches);