//! MLS group state machine.
//!
//! # Design
//!
//! [`GroupMember`] wraps an openmls [`MlsGroup`] plus the per-client
//! [`HybridCryptoProvider`] backend. The backend is **persistent** — it holds
//! the in-memory key store that maps init-key references to HPKE private keys.
//! openmls's `new_from_welcome` reads those private keys from the key store to
//! decrypt the Welcome, so the same backend instance must be used from
//! `generate_key_package` through `join_group`.
//!
//! # Hybrid post-quantum mode
//!
//! When `hybrid = true`, the backend's `derive_hpke_keypair` produces hybrid
//! (X25519 + ML-KEM-768) init keys. KeyPackages from hybrid groups contain
//! 1216-byte public keys instead of 32-byte X25519 keys. Both sender and
//! receiver must use hybrid mode for the same group.
//!
//! # Wire format
//!
//! All MLS messages are serialised/deserialised using TLS presentation language
//! encoding (`tls_codec`). The resulting byte vectors are what the transport
//! layer (and the Delivery Service) sees.
//!
//! # MLS ciphersuite
//!
//! `MLS_128_DHKEMX25519_AES128GCM_SHA256_Ed25519` — same as M2.
//!
//! # Ratchet tree
//!
//! `use_ratchet_tree_extension = true` so that the ratchet tree is embedded
//! in Welcome messages. `new_from_welcome` is called with `ratchet_tree = None`;
//! openmls extracts the tree from the Welcome's `GroupInfo` extension.

use std::{path::Path, sync::Arc};

use openmls::prelude::{
    Ciphersuite, Credential, CredentialType, CredentialWithKey, CryptoConfig, GroupId, KeyPackage,
    KeyPackageIn, MlsGroup, MlsGroupConfig, MlsMessageInBody, MlsMessageOut,
    ProcessedMessageContent, ProtocolMessage, ProtocolVersion, TlsDeserializeTrait,
    TlsSerializeTrait,
};
use openmls_traits::OpenMlsCryptoProvider;

use crate::{
    error::CoreError,
    hybrid_crypto::HybridCryptoProvider,
    identity::IdentityKeypair,
    keystore::DiskKeyStore,
};

// ── Constants ─────────────────────────────────────────────────────────────────

const CIPHERSUITE: Ciphersuite = Ciphersuite::MLS_128_DHKEMX25519_AES128GCM_SHA256_Ed25519;

// ── Return types ─────────────────────────────────────────────────────────────

/// Result of processing an incoming MLS message.
#[derive(Debug)]
pub enum ReceivedMessage {
    /// Decrypted application-layer plaintext.
    Application(Vec<u8>),
    /// Group state changed (a Commit was merged, epoch advanced).
    StateChanged,
    /// We were removed from the group by another member.
    SelfRemoved,
}

/// Like [`ReceivedMessage`] but includes the sender's credential identity bytes.
#[derive(Debug)]
pub enum ReceivedMessageWithSender {
    /// `(sender_identity, plaintext)`.
    Application(Vec<u8>, Vec<u8>),
    /// Group state changed (a Commit was merged).
    StateChanged,
    /// We were removed from the group.
    SelfRemoved,
}

// ── GroupMember ───────────────────────────────────────────────────────────────

/// Per-client MLS state: identity keypair, crypto backend, and optional group.
///
/// # Lifecycle
///
/// ```text
/// GroupMember::new(identity)
///   ├─ generate_key_package()   → upload to AS
///   ├─ create_group(group_id)   → become sole member
///   │     └─ add_member(kp)     → invite a peer; returns (commit, welcome)
///   └─ join_group(welcome)      → join after receiving a Welcome
///         ├─ send_message(msg)  → encrypt application data
///         └─ receive_message(b) → decrypt; returns Some(plaintext) or None
/// ```
pub struct GroupMember {
    /// Persistent crypto backend (hybrid or classical). Holds the in-memory key
    /// store with HPKE private keys created during `generate_key_package`.
    backend: HybridCryptoProvider,
    /// Long-term Ed25519 identity keypair. Also used as the MLS `Signer`.
    identity: Arc<IdentityKeypair>,
    /// Active MLS group, if any.
    group: Option<MlsGroup>,
    /// Shared group configuration (wire format, ratchet tree extension, etc.).
    config: MlsGroupConfig,
    /// Whether this member uses hybrid (X25519 + ML-KEM-768) HPKE keys.
    hybrid: bool,
}

impl GroupMember {
    /// Create a new `GroupMember` with a fresh classical crypto backend.
    pub fn new(identity: Arc<IdentityKeypair>) -> Self {
        Self::new_with_state(identity, DiskKeyStore::ephemeral(), None, false)
    }

    /// Create a new `GroupMember` with hybrid post-quantum crypto backend.
    pub fn new_hybrid(identity: Arc<IdentityKeypair>) -> Self {
        Self::new_with_state(identity, DiskKeyStore::ephemeral(), None, true)
    }

    /// Create a `GroupMember` with a persistent keystore at `path`.
    pub fn new_persistent(
        identity: Arc<IdentityKeypair>,
        path: impl AsRef<Path>,
    ) -> Result<Self, CoreError> {
        let key_store = DiskKeyStore::persistent(path)
            .map_err(|e| CoreError::Io(format!("keystore: {e}")))?;
        Ok(Self::new_with_state(identity, key_store, None, false))
    }

    /// Create a `GroupMember` from pre-existing state (identity + optional group + store).
    ///
    /// When `hybrid` is `true`, the backend uses hybrid (X25519 + ML-KEM-768)
    /// keys for HPKE operations. When `false`, standard X25519 keys are used.
    pub fn new_with_state(
        identity: Arc<IdentityKeypair>,
        key_store: DiskKeyStore,
        group: Option<MlsGroup>,
        hybrid: bool,
    ) -> Self {
        let config = MlsGroupConfig::builder()
            .use_ratchet_tree_extension(true)
            .build();

        let backend = if hybrid {
            HybridCryptoProvider::new_hybrid(key_store)
        } else {
            HybridCryptoProvider::new_classical(key_store)
        };

        Self {
            backend,
            identity,
            group,
            config,
            hybrid,
        }
    }

    // ── KeyPackage ────────────────────────────────────────────────────────────

    /// Generate a fresh single-use MLS KeyPackage.
    ///
    /// The HPKE init private key is stored in `self.backend`'s key store.
    /// **The same `GroupMember` instance must later call `join_group`** so
    /// that `new_from_welcome` can retrieve the private key.
    ///
    /// # Returns
    ///
    /// TLS-encoded KeyPackage bytes, ready for upload to the Authentication
    /// Service.
    ///
    /// # Errors
    ///
    /// Returns [`CoreError::Mls`] if openmls fails to create the KeyPackage.
    pub fn generate_key_package(&mut self) -> Result<Vec<u8>, CoreError> {
        let credential_with_key = self.make_credential_with_key()?;

        let key_package = KeyPackage::builder()
            .build(
                CryptoConfig::with_default_version(CIPHERSUITE),
                &self.backend,
                self.identity.as_ref(),
                credential_with_key,
            )
            .map_err(|e| CoreError::Mls(format!("{e:?}")))?;

        key_package
            .tls_serialize_detached()
            .map_err(|e| CoreError::Mls(format!("{e:?}")))
    }

    // ── Group creation ────────────────────────────────────────────────────────

    /// Create a new MLS group with `group_id` as the group identifier.
    ///
    /// The caller becomes the sole member (epoch 0). Use `add_member` to
    /// invite additional members.
    ///
    /// `group_id` can be any non-empty byte string; SHA-256 of a human-readable
    /// name is a good choice.
    ///
    /// # Errors
    ///
    /// Returns [`CoreError::Mls`] if the group already exists or openmls fails.
    pub fn create_group(&mut self, group_id: &[u8]) -> Result<(), CoreError> {
        let credential_with_key = self.make_credential_with_key()?;
        let mls_id = GroupId::from_slice(group_id);

        let group = MlsGroup::new_with_group_id(
            &self.backend,
            self.identity.as_ref(),
            &self.config,
            mls_id,
            credential_with_key,
        )
        .map_err(|e| CoreError::Mls(format!("{e:?}")))?;

        self.group = Some(group);
        Ok(())
    }

    // ── Membership ────────────────────────────────────────────────────────────

    /// Add a new member by their TLS-encoded KeyPackage bytes.
    ///
    /// Produces a Commit (to update existing members' state) and a Welcome
    /// (to bootstrap the new member). The caller is responsible for
    /// distributing these:
    ///
    /// - Send `commit_bytes` to all **existing** group members via the DS.
    ///   (In the 2-party case where the creator is the only member, this can
    ///   be discarded — the creator applies it locally via this method.)
    /// - Send `welcome_bytes` to the **new** member via the DS.
    ///
    /// This method also merges the pending Commit into the local group state
    /// (advancing the epoch), so the caller is immediately ready to encrypt.
    ///
    /// # Returns
    ///
    /// `(commit_bytes, welcome_bytes)` — both TLS-encoded MLS messages.
    ///
    /// # Errors
    ///
    /// Returns [`CoreError::Mls`] if the KeyPackage is malformed, no active
    /// group exists, or openmls fails.
    pub fn add_member(
        &mut self,
        mut key_package_bytes: &[u8],
    ) -> Result<(Vec<u8>, Vec<u8>), CoreError> {
        let group = self
            .group
            .as_mut()
            .ok_or_else(|| CoreError::Mls("no active group".into()))?;

        // Deserialise and validate the peer's KeyPackage. KeyPackage only derives
        // TlsSerialize; KeyPackageIn derives TlsDeserialize and provides validate()
        // which verifies the signature and returns a trusted KeyPackage.
        let key_package: KeyPackage =
            KeyPackageIn::tls_deserialize(&mut key_package_bytes)
                .map_err(|e| CoreError::Mls(format!("KeyPackage deserialise: {e:?}")))?
                .validate(self.backend.crypto(), ProtocolVersion::Mls10)
                .map_err(|e| CoreError::Mls(format!("KeyPackage validate: {e:?}")))?;

        // Create the Commit + Welcome. The third return value (GroupInfo) is for
        // external commits and is not needed here.
        let (commit_out, welcome_out, _group_info) = group
            .add_members(&self.backend, self.identity.as_ref(), &[key_package])
            .map_err(|e| CoreError::Mls(format!("add_members: {e:?}")))?;

        // Merge the pending Commit into our own state, advancing the epoch.
        group
            .merge_pending_commit(&self.backend)
            .map_err(|e| CoreError::Mls(format!("merge_pending_commit: {e:?}")))?;

        let commit_bytes = commit_out
            .to_bytes()
            .map_err(|e| CoreError::Mls(format!("commit serialise: {e:?}")))?;
        let welcome_bytes = welcome_out
            .to_bytes()
            .map_err(|e| CoreError::Mls(format!("welcome serialise: {e:?}")))?;

        Ok((commit_bytes, welcome_bytes))
    }

    /// Remove a member by their credential identity (Ed25519 public key bytes).
    ///
    /// Produces a Commit that must be sent to all remaining members.
    /// This method merges the pending Commit locally (advancing the epoch).
    ///
    /// # Returns
    ///
    /// `commit_bytes` — TLS-encoded MLS Commit message.
    pub fn remove_member(
        &mut self,
        member_identity: &[u8],
    ) -> Result<Vec<u8>, CoreError> {
        let group = self
            .group
            .as_mut()
            .ok_or_else(|| CoreError::Mls("no active group".into()))?;

        let leaf_index = group
            .members()
            .find(|m| m.credential.identity() == member_identity)
            .map(|m| m.index)
            .ok_or_else(|| CoreError::Mls("member not found in group".into()))?;

        let (commit_out, _welcome_option, _group_info) = group
            .remove_members(&self.backend, self.identity.as_ref(), &[leaf_index])
            .map_err(|e| CoreError::Mls(format!("remove_members: {e:?}")))?;

        group
            .merge_pending_commit(&self.backend)
            .map_err(|e| CoreError::Mls(format!("merge_pending_commit: {e:?}")))?;

        commit_out
            .to_bytes()
            .map_err(|e| CoreError::Mls(format!("commit serialise: {e:?}")))
    }

    /// Propose removing ourselves from the group (self-leave).
    ///
    /// Returns a proposal message that must be sent to other group members.
    /// Another member must then commit the pending proposals for the removal
    /// to take effect.
    pub fn leave_group(&mut self) -> Result<Vec<u8>, CoreError> {
        let group = self
            .group
            .as_mut()
            .ok_or_else(|| CoreError::Mls("no active group".into()))?;

        let proposal_out = group
            .leave_group(&self.backend, self.identity.as_ref())
            .map_err(|e| CoreError::Mls(format!("leave_group: {e:?}")))?;

        proposal_out
            .to_bytes()
            .map_err(|e| CoreError::Mls(format!("proposal serialise: {e:?}")))
    }

    /// Commit all pending proposals (stored via `store_pending_proposal`).
    ///
    /// Returns a Commit message and optionally a Welcome (if an Add was pending).
    pub fn commit_pending_proposals(
        &mut self,
    ) -> Result<(Vec<u8>, Option<Vec<u8>>), CoreError> {
        let group = self
            .group
            .as_mut()
            .ok_or_else(|| CoreError::Mls("no active group".into()))?;

        let (commit_out, welcome_option, _group_info) = group
            .commit_to_pending_proposals(&self.backend, self.identity.as_ref())
            .map_err(|e| CoreError::Mls(format!("commit_to_pending_proposals: {e:?}")))?;

        group
            .merge_pending_commit(&self.backend)
            .map_err(|e| CoreError::Mls(format!("merge_pending_commit: {e:?}")))?;

        let commit_bytes = commit_out
            .to_bytes()
            .map_err(|e| CoreError::Mls(format!("commit serialise: {e:?}")))?;

        let welcome_bytes = welcome_option
            .map(|w| {
                w.to_bytes()
                    .map_err(|e| CoreError::Mls(format!("welcome serialise: {e:?}")))
            })
            .transpose()?;

        Ok((commit_bytes, welcome_bytes))
    }

    /// Propose a self-update (key rotation).
    ///
    /// Returns a Proposal message to send to all group members.
    /// Another member (or self) must then commit the pending proposals.
    pub fn propose_self_update(&mut self) -> Result<Vec<u8>, CoreError> {
        let group = self
            .group
            .as_mut()
            .ok_or_else(|| CoreError::Mls("no active group".into()))?;

        let (proposal_out, _ref) = group
            .propose_self_update(&self.backend, self.identity.as_ref(), None)
            .map_err(|e| CoreError::Mls(format!("propose_self_update: {e:?}")))?;

        proposal_out
            .to_bytes()
            .map_err(|e| CoreError::Mls(format!("proposal serialise: {e:?}")))
    }

    /// Whether there are pending proposals waiting to be committed.
    pub fn has_pending_proposals(&self) -> bool {
        self.group
            .as_ref()
            .map(|g| g.pending_proposals().next().is_some())
            .unwrap_or(false)
    }

    /// Join an existing MLS group from a TLS-encoded Welcome message.
    ///
    /// The caller must have previously called [`generate_key_package`] on
    /// **this same instance** so that the HPKE init private key is in the
    /// backend's key store.
    ///
    /// # Errors
    ///
    /// Returns [`CoreError::Mls`] if the Welcome does not match any known
    /// KeyPackage, or openmls validation fails.
    ///
    /// [`generate_key_package`]: Self::generate_key_package
    pub fn join_group(&mut self, mut welcome_bytes: &[u8]) -> Result<(), CoreError> {
        // Deserialise MlsMessageIn, then extract the inner Welcome.
        let msg_in = openmls::prelude::MlsMessageIn::tls_deserialize(&mut welcome_bytes)
            .map_err(|e| CoreError::Mls(format!("Welcome deserialise: {e:?}")))?;

        // into_welcome() is feature-gated in openmls 0.5; extract() is public.
        let welcome = match msg_in.extract() {
            MlsMessageInBody::Welcome(w) => w,
            _ => return Err(CoreError::Mls("expected a Welcome message".into())),
        };

        // ratchet_tree = None because use_ratchet_tree_extension = true embeds
        // the tree inside the Welcome's GroupInfo extension.
        let group = MlsGroup::new_from_welcome(&self.backend, &self.config, welcome, None)
            .map_err(|e| CoreError::Mls(format!("new_from_welcome: {e:?}")))?;

        self.group = Some(group);
        Ok(())
    }

    // ── Application messages ──────────────────────────────────────────────────

    /// Encrypt `plaintext` as an MLS Application message.
    ///
    /// # Returns
    ///
    /// TLS-encoded `MlsMessageOut` bytes (PrivateMessage variant).
    ///
    /// # Errors
    ///
    /// Returns [`CoreError::Mls`] if there is no active group or encryption fails.
    pub fn send_message(&mut self, plaintext: &[u8]) -> Result<Vec<u8>, CoreError> {
        let group = self
            .group
            .as_mut()
            .ok_or_else(|| CoreError::Mls("no active group".into()))?;

        let mls_msg: MlsMessageOut = group
            .create_message(&self.backend, self.identity.as_ref(), plaintext)
            .map_err(|e| CoreError::Mls(format!("create_message: {e:?}")))?;

        mls_msg
            .to_bytes()
            .map_err(|e| CoreError::Mls(format!("message serialise: {e:?}")))
    }

    /// Process an incoming TLS-encoded MLS message.
    ///
    /// # Returns
    ///
    /// - [`ReceivedMessage::Application`] for decrypted application messages.
    /// - [`ReceivedMessage::StateChanged`] for Commits and Proposals (group state updated).
    /// - [`ReceivedMessage::SelfRemoved`] if the Commit removed us from the group.
    ///
    /// # Errors
    ///
    /// Returns [`CoreError::Mls`] if the message is malformed, fails
    /// authentication, or the group state is inconsistent.
    pub fn receive_message(&mut self, mut bytes: &[u8]) -> Result<ReceivedMessage, CoreError> {
        let group = self
            .group
            .as_mut()
            .ok_or_else(|| CoreError::Mls("no active group".into()))?;

        let msg_in = openmls::prelude::MlsMessageIn::tls_deserialize(&mut bytes)
            .map_err(|e| CoreError::Mls(format!("message deserialise: {e:?}")))?;

        // into_protocol_message() is feature-gated; extract() + manual construction is not.
        let protocol_message = match msg_in.extract() {
            MlsMessageInBody::PrivateMessage(m) => ProtocolMessage::PrivateMessage(m),
            MlsMessageInBody::PublicMessage(m) => ProtocolMessage::PublicMessage(m),
            _ => return Err(CoreError::Mls("not a protocol message".into())),
        };

        let processed = group
            .process_message(&self.backend, protocol_message)
            .map_err(|e| CoreError::Mls(format!("process_message: {e:?}")))?;

        match processed.into_content() {
            ProcessedMessageContent::ApplicationMessage(app) => {
                Ok(ReceivedMessage::Application(app.into_bytes()))
            }
            ProcessedMessageContent::StagedCommitMessage(staged) => {
                // Check if this commit removes us.
                let own_index = group.own_leaf_index();
                let self_removed = staged.remove_proposals().any(|queued| {
                    queued.remove_proposal().removed() == own_index
                });

                group
                    .merge_staged_commit(&self.backend, *staged)
                    .map_err(|e| CoreError::Mls(format!("merge_staged_commit: {e:?}")))?;

                if self_removed {
                    self.group = None;
                    Ok(ReceivedMessage::SelfRemoved)
                } else {
                    Ok(ReceivedMessage::StateChanged)
                }
            }
            // Proposals are stored for a later Commit; nothing to return yet.
            ProcessedMessageContent::ProposalMessage(proposal) => {
                group.store_pending_proposal(*proposal);
                Ok(ReceivedMessage::StateChanged)
            }
            ProcessedMessageContent::ExternalJoinProposalMessage(proposal) => {
                group.store_pending_proposal(*proposal);
                Ok(ReceivedMessage::StateChanged)
            }
        }
    }

    /// Process an incoming TLS-encoded MLS message and return sender identity + plaintext for application messages.
    ///
    /// Same as [`receive_message`], but for Application messages returns
    /// `(sender_identity_bytes, plaintext)` so the client can display who sent the message.
    pub fn receive_message_with_sender(
        &mut self,
        mut bytes: &[u8],
    ) -> Result<ReceivedMessageWithSender, CoreError> {
        let group = self
            .group
            .as_mut()
            .ok_or_else(|| CoreError::Mls("no active group".into()))?;

        let msg_in = openmls::prelude::MlsMessageIn::tls_deserialize(&mut bytes)
            .map_err(|e| CoreError::Mls(format!("message deserialise: {e:?}")))?;

        let protocol_message = match msg_in.extract() {
            MlsMessageInBody::PrivateMessage(m) => ProtocolMessage::PrivateMessage(m),
            MlsMessageInBody::PublicMessage(m) => ProtocolMessage::PublicMessage(m),
            _ => return Err(CoreError::Mls("not a protocol message".into())),
        };

        let processed = group
            .process_message(&self.backend, protocol_message)
            .map_err(|e| CoreError::Mls(format!("process_message: {e:?}")))?;

        let sender_identity = processed.credential().identity().to_vec();

        match processed.into_content() {
            ProcessedMessageContent::ApplicationMessage(app) => {
                Ok(ReceivedMessageWithSender::Application(sender_identity, app.into_bytes()))
            }
            ProcessedMessageContent::StagedCommitMessage(staged) => {
                let own_index = group.own_leaf_index();
                let self_removed = staged.remove_proposals().any(|queued| {
                    queued.remove_proposal().removed() == own_index
                });

                group
                    .merge_staged_commit(&self.backend, *staged)
                    .map_err(|e| CoreError::Mls(format!("merge_staged_commit: {e:?}")))?;

                if self_removed {
                    self.group = None;
                    Ok(ReceivedMessageWithSender::SelfRemoved)
                } else {
                    Ok(ReceivedMessageWithSender::StateChanged)
                }
            }
            ProcessedMessageContent::ProposalMessage(proposal) => {
                group.store_pending_proposal(*proposal);
                Ok(ReceivedMessageWithSender::StateChanged)
            }
            ProcessedMessageContent::ExternalJoinProposalMessage(proposal) => {
                group.store_pending_proposal(*proposal);
                Ok(ReceivedMessageWithSender::StateChanged)
            }
        }
    }

    // ── Accessors ─────────────────────────────────────────────────────────────

    /// Return the MLS group ID bytes, or `None` if no group is active.
    pub fn group_id(&self) -> Option<Vec<u8>> {
        self.group
            .as_ref()
            .map(|g| g.group_id().as_slice().to_vec())
    }

    /// Return a reference to the identity keypair.
    pub fn identity(&self) -> &IdentityKeypair {
        &self.identity
    }

    /// Return the private seed of the identity (for persistence).
    pub fn identity_seed(&self) -> [u8; 32] {
        self.identity.seed_bytes()
    }

    /// Return a reference to the underlying crypto backend.
    pub fn backend(&self) -> &HybridCryptoProvider {
        &self.backend
    }

    /// Whether this member uses hybrid post-quantum HPKE keys.
    pub fn is_hybrid(&self) -> bool {
        self.hybrid
    }

    /// Return the current MLS epoch, or `None` if no group is active.
    pub fn epoch(&self) -> Option<u64> {
        self.group.as_ref().map(|g| g.epoch().as_u64())
    }

    /// Return a reference to the MLS group, if active.
    pub fn group_ref(&self) -> Option<&MlsGroup> {
        self.group.as_ref()
    }

    /// Return the identity (credential) bytes of all current group members.
    ///
    /// Each entry is the raw credential payload (Ed25519 public key bytes)
    /// extracted from the member's MLS leaf node.
    pub fn member_identities(&self) -> Vec<Vec<u8>> {
        let group = match self.group.as_ref() {
            Some(g) => g,
            None => return Vec::new(),
        };
        group
            .members()
            .map(|m| m.credential.identity().to_vec())
            .collect()
    }

    // ── Private helpers ───────────────────────────────────────────────────────

    fn make_credential_with_key(&self) -> Result<CredentialWithKey, CoreError> {
        let credential = Credential::new(
            self.identity.public_key_bytes().to_vec(),
            CredentialType::Basic,
        )
        .map_err(|e| CoreError::Mls(format!("{e:?}")))?;

        Ok(CredentialWithKey {
            credential,
            signature_key: self.identity.public_key_bytes().to_vec().into(),
        })
    }
}

// ── Unit tests ────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    /// Full two-party MLS round-trip: creator creates group, adds joiner, then they exchange messages.
    #[test]
    fn two_party_mls_round_trip() {
        let creator_id = Arc::new(IdentityKeypair::generate());
        let joiner_id = Arc::new(IdentityKeypair::generate());

        let mut creator = GroupMember::new(Arc::clone(&creator_id));
        let mut joiner = GroupMember::new(Arc::clone(&joiner_id));

        let joiner_kp = joiner
            .generate_key_package()
            .expect("joiner KeyPackage");

        creator
            .create_group(b"test-group-m3")
            .expect("creator create group");

        let (_, welcome) = creator
            .add_member(&joiner_kp)
            .expect("creator add joiner");

        joiner.join_group(&welcome).expect("joiner join group");

        let ct_creator = creator.send_message(b"hello").expect("creator send");
        let pt_joiner = match joiner.receive_message(&ct_creator).expect("joiner recv") {
            ReceivedMessage::Application(pt) => pt,
            other => panic!("expected Application, got {other:?}"),
        };
        assert_eq!(pt_joiner, b"hello");

        let ct_joiner = joiner.send_message(b"hello back").expect("joiner send");
        let pt_creator = match creator.receive_message(&ct_joiner).expect("creator recv") {
            ReceivedMessage::Application(pt) => pt,
            other => panic!("expected Application, got {other:?}"),
        };
        assert_eq!(pt_creator, b"hello back");
    }

    /// Full two-party hybrid MLS round-trip with post-quantum HPKE keys.
    #[test]
    fn two_party_hybrid_mls_round_trip() {
        let creator_id = Arc::new(IdentityKeypair::generate());
        let joiner_id = Arc::new(IdentityKeypair::generate());

        let mut creator = GroupMember::new_hybrid(Arc::clone(&creator_id));
        let mut joiner = GroupMember::new_hybrid(Arc::clone(&joiner_id));

        assert!(creator.is_hybrid());
        assert!(joiner.is_hybrid());

        let joiner_kp = joiner
            .generate_key_package()
            .expect("joiner hybrid KeyPackage");

        creator
            .create_group(b"test-hybrid-group")
            .expect("creator create hybrid group");

        let (_, welcome) = creator
            .add_member(&joiner_kp)
            .expect("creator add joiner with hybrid KP");

        joiner.join_group(&welcome).expect("joiner join hybrid group");

        let ct_creator = creator.send_message(b"hello PQ").expect("creator send");
        let pt_joiner = match joiner.receive_message(&ct_creator).expect("joiner recv") {
            ReceivedMessage::Application(pt) => pt,
            other => panic!("expected Application, got {other:?}"),
        };
        assert_eq!(pt_joiner, b"hello PQ");

        let ct_joiner = joiner.send_message(b"quantum safe!").expect("joiner send");
        let pt_creator = match creator.receive_message(&ct_joiner).expect("creator recv") {
            ReceivedMessage::Application(pt) => pt,
            other => panic!("expected Application, got {other:?}"),
        };
        assert_eq!(pt_creator, b"quantum safe!");
    }

    /// `group_id()` returns None before create_group, Some afterwards.
    #[test]
    fn group_id_lifecycle() {
        let id = Arc::new(IdentityKeypair::generate());
        let mut member = GroupMember::new(id);

        assert!(member.group_id().is_none(), "no group before create");
        member.create_group(b"gid").unwrap();
        assert_eq!(
            member.group_id().unwrap(),
            b"gid".as_slice(),
            "group_id must match what was passed"
        );
    }

    /// Helper: set up a 3-party group (creator + A + B).
    fn setup_three_party(hybrid: bool) -> (GroupMember, GroupMember, GroupMember) {
        let creator_id = Arc::new(IdentityKeypair::generate());
        let a_id = Arc::new(IdentityKeypair::generate());
        let b_id = Arc::new(IdentityKeypair::generate());

        let (mut creator, mut a, mut b) = if hybrid {
            (
                GroupMember::new_hybrid(creator_id),
                GroupMember::new_hybrid(a_id),
                GroupMember::new_hybrid(b_id),
            )
        } else {
            (
                GroupMember::new(creator_id),
                GroupMember::new(a_id),
                GroupMember::new(b_id),
            )
        };

        let a_kp = a.generate_key_package().expect("A KeyPackage");
        let b_kp = b.generate_key_package().expect("B KeyPackage");

        creator.create_group(b"three-party").expect("create group");

        // Add A
        let (_commit_a, welcome_a) = creator.add_member(&a_kp).expect("add A");
        a.join_group(&welcome_a).expect("A join");

        // A must process the commit that added them (it's a StateChanged for A since
        // the commit itself is what brought them in — but actually A joined via Welcome,
        // so A doesn't process the add-commit). The creator already merged the pending
        // commit in add_member, so creator is at epoch 2.

        // Add B — at this point creator is at epoch 2 (after adding A).
        let (commit_b, welcome_b) = creator.add_member(&b_kp).expect("add B");
        b.join_group(&welcome_b).expect("B join");

        // A must process the commit that added B to stay in sync.
        match a.receive_message(&commit_b).expect("A recv add-B commit") {
            ReceivedMessage::StateChanged => {}
            other => panic!("expected StateChanged, got {other:?}"),
        }

        (creator, a, b)
    }

    /// Three-party hybrid MLS round-trip: all members exchange messages.
    #[test]
    fn three_party_hybrid_mls_round_trip() {
        let (mut creator, mut a, mut b) = setup_three_party(true);

        // Creator sends to A and B
        let ct = creator.send_message(b"hello group").expect("creator send");
        let pt_a = match a.receive_message(&ct).expect("A recv") {
            ReceivedMessage::Application(pt) => pt,
            other => panic!("expected Application, got {other:?}"),
        };
        let pt_b = match b.receive_message(&ct).expect("B recv") {
            ReceivedMessage::Application(pt) => pt,
            other => panic!("expected Application, got {other:?}"),
        };
        assert_eq!(pt_a, b"hello group");
        assert_eq!(pt_b, b"hello group");

        // A sends, creator and B receive
        let ct_a = a.send_message(b"from A").expect("A send");
        let pt_creator = match creator.receive_message(&ct_a).expect("creator recv") {
            ReceivedMessage::Application(pt) => pt,
            other => panic!("expected Application, got {other:?}"),
        };
        let pt_b2 = match b.receive_message(&ct_a).expect("B recv A") {
            ReceivedMessage::Application(pt) => pt,
            other => panic!("expected Application, got {other:?}"),
        };
        assert_eq!(pt_creator, b"from A");
        assert_eq!(pt_b2, b"from A");
    }

    /// Creator adds A and B, then removes B. A and creator can still communicate.
    /// B can no longer decrypt.
    #[test]
    fn three_party_remove_member() {
        let (mut creator, mut a, mut b) = setup_three_party(false);

        // Get B's identity for removal
        let b_identity = b.identity.public_key_bytes().to_vec();

        // Creator removes B
        let remove_commit = creator.remove_member(&b_identity).expect("remove B");

        // A processes the remove commit
        match a.receive_message(&remove_commit).expect("A recv remove") {
            ReceivedMessage::StateChanged => {}
            other => panic!("expected StateChanged, got {other:?}"),
        }

        // B processes the remove commit — should get SelfRemoved
        match b.receive_message(&remove_commit).expect("B recv remove") {
            ReceivedMessage::SelfRemoved => {}
            other => panic!("expected SelfRemoved, got {other:?}"),
        }

        // B's group should be cleared
        assert!(b.group_id().is_none(), "B's group should be None after removal");

        // Creator and A can still communicate
        let ct = creator.send_message(b"after removal").expect("creator send");
        let pt = match a.receive_message(&ct).expect("A recv") {
            ReceivedMessage::Application(pt) => pt,
            other => panic!("expected Application, got {other:?}"),
        };
        assert_eq!(pt, b"after removal");

        // B cannot send (no group)
        assert!(b.send_message(b"should fail").is_err());
    }

    /// A proposes to leave, creator commits the proposal, A receives SelfRemoved.
    #[test]
    fn leave_group_proposal() {
        let (mut creator, mut a, _b) = setup_three_party(false);

        // A proposes to leave
        let leave_proposal = a.leave_group().expect("A leave");

        // Creator receives the proposal (stored as pending)
        match creator.receive_message(&leave_proposal).expect("creator recv proposal") {
            ReceivedMessage::StateChanged => {}
            other => panic!("expected StateChanged for proposal, got {other:?}"),
        }

        // Creator should have pending proposals
        assert!(creator.has_pending_proposals(), "should have pending proposal");

        // Creator commits the pending proposals
        let (commit_bytes, _welcome) = creator
            .commit_pending_proposals()
            .expect("commit pending");

        // A processes the commit — should get SelfRemoved
        match a.receive_message(&commit_bytes).expect("A recv commit") {
            ReceivedMessage::SelfRemoved => {}
            other => panic!("expected SelfRemoved, got {other:?}"),
        }

        assert!(a.group_id().is_none(), "A's group should be None after leave");
    }

    /// Propose self-update, commit, other member processes the commit.
    #[test]
    fn propose_self_update_round_trip() {
        let creator_id = Arc::new(IdentityKeypair::generate());
        let joiner_id = Arc::new(IdentityKeypair::generate());

        let mut creator = GroupMember::new(Arc::clone(&creator_id));
        let mut joiner = GroupMember::new(Arc::clone(&joiner_id));

        let joiner_kp = joiner.generate_key_package().expect("joiner KP");
        creator.create_group(b"update-test").expect("create");
        let (_commit, welcome) = creator.add_member(&joiner_kp).expect("add");
        joiner.join_group(&welcome).expect("join");

        // Creator proposes a self-update
        let update_proposal = creator.propose_self_update().expect("propose update");

        // Joiner receives the proposal
        match joiner.receive_message(&update_proposal).expect("joiner recv proposal") {
            ReceivedMessage::StateChanged => {}
            other => panic!("expected StateChanged, got {other:?}"),
        }

        // Joiner commits the pending update proposal
        let (commit_bytes, _) = joiner.commit_pending_proposals().expect("commit update");

        // Creator processes the commit
        match creator.receive_message(&commit_bytes).expect("creator recv commit") {
            ReceivedMessage::StateChanged => {}
            other => panic!("expected StateChanged, got {other:?}"),
        }

        // Both can still communicate after the update
        let ct = creator.send_message(b"post-update").expect("send");
        let pt = match joiner.receive_message(&ct).expect("recv") {
            ReceivedMessage::Application(pt) => pt,
            other => panic!("expected Application, got {other:?}"),
        };
        assert_eq!(pt, b"post-update");
    }

    /// Receiving a ciphertext from a stale (lower) epoch returns an error — not a panic.
    /// This is the core invariant violated by the bidirectional-/dm race condition.
    #[test]
    fn receive_stale_epoch_message_returns_error() {
        let creator_id = Arc::new(IdentityKeypair::generate());
        let joiner_a_id = Arc::new(IdentityKeypair::generate());
        let joiner_b_id = Arc::new(IdentityKeypair::generate());

        let mut creator = GroupMember::new(Arc::clone(&creator_id));
        let mut joiner_a = GroupMember::new(Arc::clone(&joiner_a_id));
        let mut joiner_b = GroupMember::new(Arc::clone(&joiner_b_id));

        // Set up group with joiner_a (epoch 1 after create_group, epoch 2 after add).
        let kp_a = joiner_a.generate_key_package().expect("kp_a");
        creator.create_group(b"stale-epoch-test").expect("create");
        let (_, welcome_a) = creator.add_member(&kp_a).expect("add a");
        joiner_a.join_group(&welcome_a).expect("join a");

        // Creator sends a message at the current epoch (epoch 2).
        let ct_epoch2 = creator.send_message(b"epoch-2 message").expect("send");

        // Creator now adds joiner_b, advancing to epoch 3. joiner_a must process the commit.
        let kp_b = joiner_b.generate_key_package().expect("kp_b");
        let (commit_b, welcome_b) = creator.add_member(&kp_b).expect("add b");
        joiner_b.join_group(&welcome_b).expect("join b");
        match joiner_a.receive_message(&commit_b).expect("a recv add-b commit") {
            ReceivedMessage::StateChanged => {}
            other => panic!("expected StateChanged, got {other:?}"),
        }

        // joiner_b joined at epoch 3 via Welcome. Attempting to decrypt ct_epoch2 (epoch 2)
        // must return an error, not panic.
        let result = joiner_b.receive_message(&ct_epoch2);
        assert!(
            result.is_err(),
            "decrypting an epoch-2 ciphertext in epoch-3 context must fail, not panic"
        );
    }

    /// 10 messages alternating Alice→Bob and Bob→Alice all decrypt successfully.
    /// Verifies that epoch state stays in sync across multiple application messages.
    #[test]
    fn multi_message_roundtrip_epoch_stays_in_sync() {
        let alice_id = Arc::new(IdentityKeypair::generate());
        let bob_id = Arc::new(IdentityKeypair::generate());

        let mut alice = GroupMember::new(Arc::clone(&alice_id));
        let mut bob = GroupMember::new(Arc::clone(&bob_id));

        let bob_kp = bob.generate_key_package().expect("bob kp");
        alice.create_group(b"multi-msg-test").expect("create");
        let (_, welcome) = alice.add_member(&bob_kp).expect("add bob");
        bob.join_group(&welcome).expect("join");

        for i in 0u32..5 {
            let payload_alice = format!("alice msg {i}");
            let ct = alice.send_message(payload_alice.as_bytes()).expect("alice send");
            let pt = match bob.receive_message(&ct).expect("bob recv") {
                ReceivedMessage::Application(pt) => pt,
                other => panic!("expected Application, got {other:?}"),
            };
            assert_eq!(pt, payload_alice.as_bytes());

            let payload_bob = format!("bob reply {i}");
            let ct = bob.send_message(payload_bob.as_bytes()).expect("bob send");
            let pt = match alice.receive_message(&ct).expect("alice recv") {
                ReceivedMessage::Application(pt) => pt,
                other => panic!("expected Application, got {other:?}"),
            };
            assert_eq!(pt, payload_bob.as_bytes());
        }
    }

    /// A member who has not yet joined (no group) cannot send messages.
    #[test]
    fn send_before_join_returns_error() {
        let id = Arc::new(IdentityKeypair::generate());
        let mut member = GroupMember::new(id);
        assert!(
            member.send_message(b"too early").is_err(),
            "send_message before join must return an error"
        );
    }
}