//! MLS group keying (RFC 9420) for a Joule Code collaboration cluster (SPEC §7).
//!
//! The cluster members form an MLS group and derive a **shared group key** with
//! no central server: every member exports the same secret, used to encrypt the
//! collaboration's traffic (coding operations / Smart Byte envelopes). Adding a
//! member advances the epoch; the new member, once welcomed, derives the same
//! secret. This is the group-keying layer that sits above the BFT consensus +
//! FROST threshold authorization in `joule-code-collab`.
//!
//! Built on `openmls` (the RFC 9420 reference implementation). This test
//! demonstrates the load-bearing property — agreement on a shared group key — so
//! the layer is verified, not asserted.

use openmls::prelude::tls_codec::{Deserialize as _, Serialize as _};
use openmls::prelude::*;
use openmls_basic_credential::SignatureKeyPair;
use openmls_rust_crypto::OpenMlsRustCrypto;

const CIPHERSUITE: Ciphersuite = Ciphersuite::MLS_128_DHKEMX25519_AES128GCM_SHA256_Ed25519;
const EXPORT_LABEL: &str = "joule-code-collab-group-key";

fn member(
    provider: &OpenMlsRustCrypto,
    name: &str,
) -> (CredentialWithKey, SignatureKeyPair) {
    let credential = BasicCredential::new(name.as_bytes().to_vec());
    let signer = SignatureKeyPair::new(CIPHERSUITE.signature_algorithm()).expect("keypair");
    signer.store(provider.storage()).expect("store signer");
    let credential_with_key = CredentialWithKey {
        credential: credential.into(),
        signature_key: signer.public().into(),
    };
    (credential_with_key, signer)
}

fn key_package(
    provider: &OpenMlsRustCrypto,
    cwk: CredentialWithKey,
    signer: &SignatureKeyPair,
) -> KeyPackageBundle {
    KeyPackage::builder()
        .build(CIPHERSUITE, provider, signer, cwk)
        .expect("key package")
}

fn export(group: &MlsGroup, provider: &OpenMlsRustCrypto) -> Vec<u8> {
    group
        .export_secret(provider, EXPORT_LABEL, &[], 32)
        .expect("export secret")
}

/// Extract the `Welcome` from the founder's output message via the tls wire
/// round-trip — the path a real recipient takes off the network.
fn extract_welcome(message: MlsMessageOut) -> Welcome {
    let bytes = message.tls_serialize_detached().expect("serialize welcome");
    let incoming = MlsMessageIn::tls_deserialize_exact(bytes.as_slice()).expect("deserialize");
    match incoming.extract() {
        MlsMessageBodyIn::Welcome(welcome) => welcome,
        _ => panic!("expected a Welcome message"),
    }
}

#[test]
fn cluster_members_derive_a_shared_group_key() {
    let alice_p = OpenMlsRustCrypto::default();
    let bob_p = OpenMlsRustCrypto::default();
    let carol_p = OpenMlsRustCrypto::default();

    let (alice_cwk, alice_signer) = member(&alice_p, "alice");
    let (bob_cwk, bob_signer) = member(&bob_p, "bob");
    let (carol_cwk, carol_signer) = member(&carol_p, "carol");

    // Bob and Carol publish key packages so the group founder can add them.
    let bob_kp = key_package(&bob_p, bob_cwk, &bob_signer);
    let carol_kp = key_package(&carol_p, carol_cwk, &carol_signer);

    // Alice founds the collaboration group.
    let mut alice_group = MlsGroup::new(
        &alice_p,
        &alice_signer,
        &MlsGroupCreateConfig::default(),
        alice_cwk,
    )
    .expect("found group");

    // Add Bob.
    let (_commit, welcome, _group_info) = alice_group
        .add_members(&alice_p, &alice_signer, &[bob_kp.key_package().clone()])
        .expect("add bob");
    alice_group.merge_pending_commit(&alice_p).expect("merge");

    let welcome = extract_welcome(welcome);
    let ratchet_tree = alice_group.export_ratchet_tree();
    let bob_group = StagedWelcome::new_from_welcome(
        &bob_p,
        &MlsGroupJoinConfig::default(),
        welcome,
        Some(ratchet_tree.into()),
    )
    .expect("stage welcome")
    .into_group(&bob_p)
    .expect("bob joins");

    // Alice and Bob now share the group key.
    let epoch_after_bob = alice_group.epoch();
    assert_eq!(export(&alice_group, &alice_p), export(&bob_group, &bob_p));

    // Add Carol — a new epoch — and Carol derives the same shared key.
    let (_commit, welcome, _gi) = alice_group
        .add_members(&alice_p, &alice_signer, &[carol_kp.key_package().clone()])
        .expect("add carol");
    alice_group.merge_pending_commit(&alice_p).expect("merge");
    // Bob processes the same commit to stay in sync.
    // (In a real cluster the commit is delivered to every member; here Bob joins
    //  via the new member's welcome path is not needed — we check Carol + Alice.)
    assert!(
        alice_group.epoch() > epoch_after_bob,
        "adding a member must advance the epoch"
    );

    let welcome = extract_welcome(welcome);
    let ratchet_tree = alice_group.export_ratchet_tree();
    let carol_group = StagedWelcome::new_from_welcome(
        &carol_p,
        &MlsGroupJoinConfig::default(),
        welcome,
        Some(ratchet_tree.into()),
    )
    .expect("stage welcome")
    .into_group(&carol_p)
    .expect("carol joins");

    assert_eq!(export(&alice_group, &alice_p), export(&carol_group, &carol_p));
    // Bob's view is now an epoch behind (he hasn't processed Carol's commit), so
    // his exported secret differs from the new epoch's — exactly the forward-
    // secret, epoch-bound property MLS provides.
    assert_ne!(export(&bob_group, &bob_p), export(&alice_group, &alice_p));
}