use std::collections::HashMap;
use std::path::Path;

use mls_rs::crypto::{SignaturePublicKey, SignatureSecretKey};

use mls_rs::group::ReceivedMessage;
use mls_rs::identity::basic::{BasicCredential, BasicIdentityProvider};
use mls_rs::identity::SigningIdentity;
use mls_rs::{CipherSuite, CipherSuiteProvider, Client, CryptoProvider, ExtensionList, MlsMessage};
use mls_rs_provider_sqlite::connection_strategy::FileConnectionStrategy;
use zeroize::Zeroizing;

use crate::zeroizing_crypto::ZeroizingCryptoProvider;
use mls_rs_provider_sqlite::storage::{
    SqLiteGroupStateStorage, SqLiteKeyPackageStorage, SqLitePreSharedKeyStorage,
};
use mls_rs_provider_sqlite::SqLiteDataStorageEngine;

const MLS_CIPHERSUITE: CipherSuite = CipherSuite::CURVE25519_AES128;

/// Identity encoding version for DID:web credentials.
const IDENTITY_VERSION_DID: u8 = 1;

/// Identity encoding version for DID:key credentials.
const IDENTITY_VERSION_DID_KEY: u8 = 2;

/// A Skytale identity that can be a raw Ed25519 public key, a DID:web
/// identifier, or a DID:key identifier, each paired with a public key.
///
/// This mirrors the `SkytaleIdentity` type from `skytale-store` but is
/// self-contained within the SDK crate to avoid a dependency on the
/// workspace crate.
///
/// ## Wire Format
///
/// - **Version 0 (raw):** Plain bytes used directly as the MLS credential
///   identity (no CBOR wrapper). This is the legacy format.
/// - **Version 1 (DID:web):** CBOR-encoded array `[1, "did:web:example.com",
///   <32-byte pubkey>]`.
/// - **Version 2 (DID:key):** CBOR-encoded array `[2, "did:key:z6Mk...",
///   <32-byte pubkey>]`.
///
/// # Example
///
/// ```
/// use skytale_sdk::mls_engine::SkytaleIdentity;
///
/// let raw = SkytaleIdentity::RawKey(b"agent-1".to_vec());
/// assert_eq!(raw.to_credential_bytes(), b"agent-1");
///
/// let did = SkytaleIdentity::DidWeb {
///     uri: "did:web:example.com".to_string(),
///     pubkey: [1u8; 32],
/// };
/// let bytes = did.to_credential_bytes();
/// assert!(bytes.len() > 32);
///
/// let did_key = SkytaleIdentity::DidKey {
///     did_uri: "did:key:z6MkTest".to_string(),
///     pubkey: [2u8; 32],
/// };
/// let bytes = did_key.to_credential_bytes();
/// assert!(bytes.len() > 32);
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SkytaleIdentity {
    /// Raw identity bytes (legacy format, no CBOR wrapper).
    RawKey(Vec<u8>),
    /// A DID:web identifier with its associated Ed25519 public key.
    DidWeb {
        /// The DID URI, e.g. `"did:web:example.com"`.
        uri: String,
        /// The 32-byte Ed25519 public key.
        pubkey: [u8; 32],
    },
    /// A DID:key identifier with its associated Ed25519 public key.
    DidKey {
        /// The DID URI, e.g. `"did:key:z6MkhaXgBZDvotDkL5257faiztiGiC2QtKLGpbnnEGta2doK"`.
        did_uri: String,
        /// The 32-byte Ed25519 public key.
        pubkey: [u8; 32],
    },
}

impl SkytaleIdentity {
    /// Serialize this identity to bytes for MLS credential embedding.
    ///
    /// - `RawKey` produces the raw bytes directly.
    /// - `DidWeb` produces CBOR-encoded `[1, uri, pubkey]`.
    /// - `DidKey` produces CBOR-encoded `[2, did_uri, pubkey]`.
    ///
    /// # Example
    ///
    /// ```
    /// use skytale_sdk::mls_engine::SkytaleIdentity;
    ///
    /// let raw = SkytaleIdentity::RawKey(b"test".to_vec());
    /// assert_eq!(raw.to_credential_bytes(), b"test");
    /// ```
    pub fn to_credential_bytes(&self) -> Vec<u8> {
        match self {
            SkytaleIdentity::RawKey(key) => key.clone(),
            SkytaleIdentity::DidWeb { uri, pubkey } => {
                let payload: (u8, &str, &[u8]) = (IDENTITY_VERSION_DID, uri.as_str(), pubkey);
                let mut buf = Vec::new();
                ciborium::into_writer(&payload, &mut buf)
                    .expect("CBOR serialization of SkytaleIdentity should not fail");
                buf
            }
            SkytaleIdentity::DidKey { did_uri, pubkey } => {
                let payload: (u8, &str, &[u8]) =
                    (IDENTITY_VERSION_DID_KEY, did_uri.as_str(), pubkey);
                let mut buf = Vec::new();
                ciborium::into_writer(&payload, &mut buf)
                    .expect("CBOR serialization of SkytaleIdentity should not fail");
                buf
            }
        }
    }

    /// Parse identity bytes from an MLS credential.
    ///
    /// Tries CBOR decoding first. If the payload is a CBOR array whose first
    /// element is a recognized version tag, the versioned format is used.
    /// Otherwise the bytes are treated as raw identity bytes.
    ///
    /// # Example
    ///
    /// ```
    /// use skytale_sdk::mls_engine::SkytaleIdentity;
    ///
    /// let raw = SkytaleIdentity::RawKey(b"test".to_vec());
    /// let bytes = raw.to_credential_bytes();
    /// let parsed = SkytaleIdentity::parse(&bytes);
    /// assert_eq!(parsed, raw);
    /// ```
    pub fn parse(bytes: &[u8]) -> Self {
        if let Ok((version, uri, pubkey_bytes)) =
            ciborium::from_reader::<(u8, String, Vec<u8>), _>(bytes)
        {
            if pubkey_bytes.len() == 32 {
                let mut pubkey = [0u8; 32];
                pubkey.copy_from_slice(&pubkey_bytes);
                if version == IDENTITY_VERSION_DID {
                    return SkytaleIdentity::DidWeb { uri, pubkey };
                }
                if version == IDENTITY_VERSION_DID_KEY {
                    return SkytaleIdentity::DidKey {
                        did_uri: uri,
                        pubkey,
                    };
                }
            }
        }

        SkytaleIdentity::RawKey(bytes.to_vec())
    }

    /// Extract the Ed25519 public key bytes regardless of identity type.
    ///
    /// For `RawKey`, returns the raw bytes (which may not be exactly 32 bytes).
    /// For `DidWeb` and `DidKey`, returns the 32-byte public key.
    ///
    /// # Example
    ///
    /// ```
    /// use skytale_sdk::mls_engine::SkytaleIdentity;
    ///
    /// let did_key = SkytaleIdentity::DidKey {
    ///     did_uri: "did:key:z6MkTest".to_string(),
    ///     pubkey: [9u8; 32],
    /// };
    /// assert_eq!(did_key.pubkey(), &[9u8; 32]);
    /// ```
    pub fn pubkey(&self) -> &[u8] {
        match self {
            SkytaleIdentity::RawKey(key) => key,
            SkytaleIdentity::DidWeb { pubkey, .. } | SkytaleIdentity::DidKey { pubkey, .. } => {
                pubkey
            }
        }
    }
}

/// Result of processing an incoming MLS message.
#[derive(Debug)]
pub enum ProcessedMessage {
    /// Decrypted application data.
    Application(Vec<u8>),
    /// Commit processed, epoch advanced.
    Commit,
    /// Other MLS message type (proposal, etc.)
    Other,
}

/// Errors from MLS operations.
#[derive(Debug, thiserror::Error)]
pub enum MlsEngineError {
    #[error("MLS client error: {0}")]
    Client(String),

    #[error("MLS group error: {0}")]
    Group(String),

    #[error("group not found for id: {0}")]
    GroupNotFound(String),

    #[error("message serialization error: {0}")]
    Serialization(String),

    #[error("storage error: {0}")]
    Storage(String),

    #[error("unexpected message type")]
    UnexpectedMessageType,
}

/// Concrete mls-rs config type after building with SQLite + RustCrypto + BasicIdentity.
type SdkMlsConfig = mls_rs::client_builder::WithCryptoProvider<
    ZeroizingCryptoProvider,
    mls_rs::client_builder::WithIdentityProvider<
        BasicIdentityProvider,
        mls_rs::client_builder::WithPskStore<
            SqLitePreSharedKeyStorage,
            mls_rs::client_builder::WithKeyPackageRepo<
                SqLiteKeyPackageStorage,
                mls_rs::client_builder::WithGroupStateStorage<
                    SqLiteGroupStateStorage,
                    mls_rs::client_builder::BaseConfig,
                >,
            >,
        >,
    >,
>;

type MlsClient = Client<SdkMlsConfig>;
type MlsGroup = mls_rs::Group<SdkMlsConfig>;

/// Wraps mls-rs with persistent SQLite storage for MLS group operations.
///
/// # Example
///
/// ```no_run
/// use skytale_sdk::mls_engine::MlsEngine;
///
/// // Raw identity (backward compatible)
/// let engine = MlsEngine::new(
///     std::path::Path::new("/tmp/mls-data"),
///     b"agent-1",
/// ).unwrap();
///
/// // DID:web identity — automatically CBOR-encoded with the signing pubkey
/// let engine2 = MlsEngine::new(
///     std::path::Path::new("/tmp/mls-data-2"),
///     b"did:web:agent.example.com",
/// ).unwrap();
///
/// // DID:key identity — self-certifying, CBOR-encoded with the signing pubkey
/// let engine3 = MlsEngine::new(
///     std::path::Path::new("/tmp/mls-data-3"),
///     b"did:key:z6MkhaXgBZDvotDkL5257faiztiGiC2QtKLGpbnnEGta2doK",
/// ).unwrap();
/// ```
pub struct MlsEngine {
    /// In-memory cache of active groups keyed by hex-encoded group_id.
    groups: HashMap<String, MlsGroup>,
    /// The mls-rs client used to create/join/load groups.
    client: MlsClient,
}

impl MlsEngine {
    /// Create a new MLS engine with state persisted to the given directory.
    ///
    /// Accepts raw identity bytes. If the bytes are valid UTF-8 starting with
    /// `"did:key:"`, the identity is encoded as a [`SkytaleIdentity::DidKey`].
    /// If they start with any other `"did:"` prefix, the identity is encoded
    /// as a [`SkytaleIdentity::DidWeb`]. Both DID variants embed the Ed25519
    /// public key from the signing keypair (generated or loaded from disk).
    /// Otherwise the bytes are used as a [`SkytaleIdentity::RawKey`].
    ///
    /// This allows Python callers to simply pass `b"did:web:agent.example.com"`
    /// or `b"did:key:z6Mk..."` and have the correct CBOR credential produced
    /// in Rust.
    pub fn new(data_dir: &Path, identity: &[u8]) -> Result<Self, MlsEngineError> {
        Self::new_with_bytes(data_dir, identity.to_vec())
    }

    /// Internal constructor that takes raw identity bytes.
    fn new_with_bytes(data_dir: &Path, identity: Vec<u8>) -> Result<Self, MlsEngineError> {
        // Check for symlink before create_dir_all to prevent a TOCTOU attack
        // where an attacker replaces the directory with a symlink.
        if data_dir.exists() {
            #[cfg(unix)]
            {
                let metadata = std::fs::symlink_metadata(data_dir).map_err(|e| {
                    MlsEngineError::Client(format!("failed to check data dir: {e}"))
                })?;
                if metadata.file_type().is_symlink() {
                    return Err(MlsEngineError::Client(
                        "data directory is a symlink, refusing to use".into(),
                    ));
                }
            }
        }

        std::fs::create_dir_all(data_dir).map_err(|e| MlsEngineError::Storage(e.to_string()))?;

        // Restrict data directory to owner-only access (rwx------).
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            std::fs::set_permissions(data_dir, std::fs::Permissions::from_mode(0o700))
                .map_err(|e| MlsEngineError::Storage(e.to_string()))?;
        }

        let db_path = data_dir.join("mls.db");
        let conn_strategy = FileConnectionStrategy::new(&db_path);
        // NOTE: mls-rs key types do not implement ZeroizeOnDrop — key material
        // in the SQLite database and in-memory mls-rs structures is not
        // automatically zeroized when dropped.
        let storage_engine = SqLiteDataStorageEngine::new(conn_strategy)
            .map_err(|e| MlsEngineError::Storage(format!("{e:?}")))?;

        // Restrict db file to owner-only read/write (rw-------).
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            if db_path.exists() {
                std::fs::set_permissions(&db_path, std::fs::Permissions::from_mode(0o600))
                    .map_err(|e| MlsEngineError::Storage(e.to_string()))?;
            }
        }

        let group_state = storage_engine
            .group_state_storage()
            .map_err(|e| MlsEngineError::Storage(format!("{e:?}")))?;
        let key_package = storage_engine
            .key_package_storage()
            .map_err(|e| MlsEngineError::Storage(format!("{e:?}")))?;
        let psk = storage_engine
            .pre_shared_key_storage()
            .map_err(|e| MlsEngineError::Storage(format!("{e:?}")))?;

        let crypto = ZeroizingCryptoProvider::default();
        let cs = crypto
            .cipher_suite_provider(MLS_CIPHERSUITE)
            .ok_or_else(|| MlsEngineError::Client("unsupported cipher suite".into()))?;

        // Persist the signing key so the agent keeps the same MLS identity
        // across restarts instead of generating a new one each time.
        let key_path = data_dir.join("signing_key.bin");
        let (secret_key, public_key) = if key_path.exists() {
            let key_bytes =
                Zeroizing::new(std::fs::read(&key_path).map_err(|e| {
                    MlsEngineError::Client(format!("failed to read signing key: {e}"))
                })?);
            deserialize_signing_key(&key_bytes)?
        } else {
            let keys = cs
                .signature_key_generate()
                .map_err(|e| MlsEngineError::Client(format!("{e:?}")))?;
            let key_bytes = serialize_signing_key(&keys);
            std::fs::write(&key_path, key_bytes.as_slice())
                .map_err(|e| MlsEngineError::Client(format!("failed to write signing key: {e}")))?;
            #[cfg(unix)]
            {
                use std::os::unix::fs::PermissionsExt;
                std::fs::set_permissions(&key_path, std::fs::Permissions::from_mode(0o600))
                    .map_err(|e| {
                        MlsEngineError::Client(format!("failed to set key permissions: {e}"))
                    })?;
            }
            keys
        };

        // Detect DID URIs: if the raw identity bytes are valid UTF-8 starting
        // with "did:", construct the appropriate DID identity variant with the
        // signing public key. "did:key:" maps to DidKey, all other "did:"
        // prefixes map to DidWeb. Otherwise use the bytes as-is (RawKey).
        let credential_bytes = match std::str::from_utf8(&identity) {
            Ok(s) if s.starts_with("did:") => {
                let mut pubkey = [0u8; 32];
                let pk_bytes: &[u8] = public_key.as_ref();
                if pk_bytes.len() == 32 {
                    pubkey.copy_from_slice(pk_bytes);
                } else {
                    return Err(MlsEngineError::Client(format!(
                        "expected 32-byte Ed25519 public key, got {} bytes",
                        pk_bytes.len()
                    )));
                }
                if s.starts_with("did:key:") {
                    SkytaleIdentity::DidKey {
                        did_uri: s.to_string(),
                        pubkey,
                    }
                    .to_credential_bytes()
                } else {
                    SkytaleIdentity::DidWeb {
                        uri: s.to_string(),
                        pubkey,
                    }
                    .to_credential_bytes()
                }
            }
            _ => identity,
        };

        let basic_cred = BasicCredential::new(credential_bytes);
        let signing_id = SigningIdentity::new(basic_cred.into_credential(), public_key);

        let client = Client::builder()
            .group_state_storage(group_state)
            .key_package_repo(key_package)
            .psk_store(psk)
            .identity_provider(BasicIdentityProvider)
            .crypto_provider(ZeroizingCryptoProvider::default())
            .signing_identity(signing_id, secret_key, MLS_CIPHERSUITE)
            .build();

        Ok(Self {
            groups: HashMap::new(),
            client,
        })
    }

    /// Create a new MLS group. Returns the group_id bytes.
    pub fn create_group(&mut self, group_id: &[u8]) -> Result<Vec<u8>, MlsEngineError> {
        let group = self
            .client
            .create_group_with_id(
                group_id.to_vec(),
                ExtensionList::default(),
                Default::default(),
                None,
            )
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?;

        let id = hex::encode(group.group_id());
        self.groups.insert(id, group);
        Ok(group_id.to_vec())
    }

    /// Generate a KeyPackage for others to add us to their group.
    pub fn generate_key_package(&self) -> Result<Vec<u8>, MlsEngineError> {
        let kp_msg = self
            .client
            .generate_key_package_message(Default::default(), Default::default(), None)
            .map_err(|e| MlsEngineError::Client(format!("{e:?}")))?;
        kp_msg
            .to_bytes()
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))
    }

    /// Add a member to a group we own. Returns `(commit_bytes, welcome_bytes)`.
    ///
    /// This is a convenience method that creates the commit and applies it
    /// immediately. Use `create_add_member_commit` + `apply_pending_commit`
    /// when the commit must only be applied after transport publish succeeds.
    pub fn add_member(
        &mut self,
        group_id: &[u8],
        key_package_bytes: &[u8],
    ) -> Result<(Vec<u8>, Vec<u8>), MlsEngineError> {
        let (commit_bytes, welcome_bytes) =
            self.create_add_member_commit(group_id, key_package_bytes)?;
        self.apply_pending_commit(group_id)?;
        Ok((commit_bytes, welcome_bytes))
    }

    /// Create an add-member commit without applying it.
    ///
    /// The commit is left pending -- call `apply_pending_commit` after the
    /// commit message has been successfully published to the transport.
    pub fn create_add_member_commit(
        &mut self,
        group_id: &[u8],
        key_package_bytes: &[u8],
    ) -> Result<(Vec<u8>, Vec<u8>), MlsEngineError> {
        let group = self.get_group_mut(group_id)?;

        let kp_msg = MlsMessage::from_bytes(key_package_bytes)
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))?;

        let commit_output = group
            .commit_builder()
            .add_member(kp_msg)
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?
            .build()
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?;

        let commit_bytes = commit_output
            .commit_message
            .to_bytes()
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))?;

        let welcome_bytes = commit_output
            .welcome_messages
            .first()
            .ok_or_else(|| MlsEngineError::Group("no welcome message produced".into()))?
            .to_bytes()
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))?;

        Ok((commit_bytes, welcome_bytes))
    }

    /// Apply a pending commit for the given group.
    ///
    /// Call this after the commit has been successfully published to the
    /// transport. If the transport publish fails, the commit remains pending
    /// and the group state is unchanged.
    pub fn apply_pending_commit(&mut self, group_id: &[u8]) -> Result<(), MlsEngineError> {
        let group = self.get_group_mut(group_id)?;
        group
            .apply_pending_commit()
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?;
        Ok(())
    }

    /// Join a group from a Welcome message. Returns the group_id.
    pub fn join_from_welcome(&mut self, welcome_bytes: &[u8]) -> Result<Vec<u8>, MlsEngineError> {
        let welcome_msg = MlsMessage::from_bytes(welcome_bytes)
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))?;

        let (group, _info) = self
            .client
            .join_group(None, &welcome_msg, None)
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?;

        let group_id = group.group_id().to_vec();
        let id = hex::encode(&group_id);
        self.groups.insert(id, group);
        Ok(group_id)
    }

    /// Encrypt a message for a group.
    #[tracing::instrument(skip(self, group_id, plaintext), fields(payload.size = plaintext.len()))]
    pub fn encrypt(
        &mut self,
        group_id: &[u8],
        plaintext: &[u8],
    ) -> Result<Vec<u8>, MlsEngineError> {
        let group = self.get_group_mut(group_id)?;
        let mls_msg = group
            .encrypt_application_message(plaintext, vec![])
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?;
        mls_msg
            .to_bytes()
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))
    }

    /// Decrypt a message from a group.
    #[tracing::instrument(skip(self, group_id, ciphertext), fields(payload.size = ciphertext.len()))]
    pub fn decrypt(
        &mut self,
        group_id: &[u8],
        ciphertext: &[u8],
    ) -> Result<Vec<u8>, MlsEngineError> {
        let group = self.get_group_mut(group_id)?;
        let mls_msg = MlsMessage::from_bytes(ciphertext)
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))?;

        match group
            .process_incoming_message(mls_msg)
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?
        {
            ReceivedMessage::ApplicationMessage(app_msg) => Ok(app_msg.data().to_vec()),
            _ => Err(MlsEngineError::UnexpectedMessageType),
        }
    }

    /// Process an incoming commit (epoch advancement).
    pub fn process_commit(
        &mut self,
        group_id: &[u8],
        commit_bytes: &[u8],
    ) -> Result<(), MlsEngineError> {
        let group = self.get_group_mut(group_id)?;
        let mls_msg = MlsMessage::from_bytes(commit_bytes)
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))?;

        match group
            .process_incoming_message(mls_msg)
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?
        {
            ReceivedMessage::Commit(_) => Ok(()),
            _ => Err(MlsEngineError::UnexpectedMessageType),
        }
    }

    /// Process any incoming MLS message, letting mls-rs determine the type.
    ///
    /// Use this instead of `decrypt`/`process_commit` when the message type is
    /// unknown (e.g., the relay labels everything as `Msg`).
    pub fn process_incoming(
        &mut self,
        group_id: &[u8],
        ciphertext: &[u8],
    ) -> Result<ProcessedMessage, MlsEngineError> {
        let group = self.get_group_mut(group_id)?;
        let mls_msg = MlsMessage::from_bytes(ciphertext)
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))?;

        match group
            .process_incoming_message(mls_msg)
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?
        {
            ReceivedMessage::ApplicationMessage(app_msg) => {
                Ok(ProcessedMessage::Application(app_msg.data().to_vec()))
            }
            ReceivedMessage::Commit(_) => Ok(ProcessedMessage::Commit),
            _ => Ok(ProcessedMessage::Other),
        }
    }

    /// Load a group from persistent storage by its ID.
    pub fn load_group(&mut self, group_id: &[u8]) -> Result<(), MlsEngineError> {
        let group = self
            .client
            .load_group(group_id)
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?;
        let id = hex::encode(group_id);
        self.groups.insert(id, group);
        Ok(())
    }

    /// Rotate the signing key for a group (MLS UpdatePath).
    ///
    /// Builds an empty commit which includes a fresh UpdatePath, effectively
    /// rotating the sender's leaf key material. The commit is left pending --
    /// call [`MlsEngine::apply_pending_commit`] after the commit message has
    /// been successfully published to the transport.
    ///
    /// Returns the serialized MLS commit message that must be sent to all
    /// group members via the relay.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # let mut engine: skytale_sdk::mls_engine::MlsEngine = todo!();
    /// let commit_bytes = engine.rotate_key(b"group-id-bytes").unwrap();
    /// // Send commit_bytes to relay for distribution, then:
    /// engine.apply_pending_commit(b"group-id-bytes").unwrap();
    /// ```
    pub fn rotate_key(&mut self, group_id: &[u8]) -> Result<Vec<u8>, MlsEngineError> {
        let group = self.get_group_mut(group_id)?;

        // An empty commit_builder().build() generates a fresh UpdatePath
        // with new leaf key material per RFC 9420.
        let commit_output = group
            .commit_builder()
            .build()
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?;

        commit_output
            .commit_message
            .to_bytes()
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))
    }

    /// Export a per-channel audit encryption key using the MLS exporter (RFC 9420 Section 8.5).
    ///
    /// Returns a 32-byte secret unique to this channel and epoch, suitable for
    /// deriving an AES-256-GCM key for encrypting audit log entries. The secret
    /// changes every epoch (i.e., after each commit), providing forward secrecy
    /// for audit data.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # let mut engine: skytale_sdk::mls_engine::MlsEngine = todo!();
    /// let audit_key = engine.export_audit_key(b"group-id-bytes").unwrap();
    /// assert_eq!(audit_key.len(), 32);
    /// ```
    pub fn export_audit_key(&mut self, group_id: &[u8]) -> Result<Vec<u8>, MlsEngineError> {
        let group = self.get_group_mut(group_id)?;

        let secret = group
            .export_secret(b"skytale-audit-v1", b"", 32)
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?;

        Ok(secret.to_vec())
    }

    /// Remove a group from the in-memory cache.
    ///
    /// Returns `true` if the group was present, `false` otherwise.
    pub fn remove_group(&mut self, channel_id: &[u8]) -> bool {
        let id = hex::encode(channel_id);
        self.groups.remove(&id).is_some()
    }

    fn get_group_mut(&mut self, group_id: &[u8]) -> Result<&mut MlsGroup, MlsEngineError> {
        let id = hex::encode(group_id);
        if !self.groups.contains_key(&id) {
            if self.groups.len() >= MAX_ACTIVE_GROUPS {
                return Err(MlsEngineError::Client("too many active groups".into()));
            }
            // Try to load from persistent storage
            if let Ok(group) = self.client.load_group(group_id) {
                self.groups.insert(id.clone(), group);
            }
        }
        self.groups
            .get_mut(&id)
            .ok_or(MlsEngineError::GroupNotFound(id))
    }
}

const MAX_ACTIVE_GROUPS: usize = 1000;

/// Serialize a signing key pair as `[secret_len as u32 LE][secret_bytes][public_bytes]`.
fn serialize_signing_key(keys: &(SignatureSecretKey, SignaturePublicKey)) -> Zeroizing<Vec<u8>> {
    let secret_bytes: &[u8] = keys.0.as_ref();
    let public_bytes: &[u8] = keys.1.as_ref();
    let mut buf = Vec::with_capacity(4 + secret_bytes.len() + public_bytes.len());
    buf.extend_from_slice(&(secret_bytes.len() as u32).to_le_bytes());
    buf.extend_from_slice(secret_bytes);
    buf.extend_from_slice(public_bytes);
    Zeroizing::new(buf)
}

/// Deserialize a signing key pair from the format produced by `serialize_signing_key`.
fn deserialize_signing_key(
    data: &[u8],
) -> Result<(SignatureSecretKey, SignaturePublicKey), MlsEngineError> {
    if data.len() < 4 {
        return Err(MlsEngineError::Client("signing key file too short".into()));
    }
    let secret_len = u32::from_le_bytes(
        data[..4]
            .try_into()
            .map_err(|_| MlsEngineError::Client("invalid key length prefix".into()))?,
    ) as usize;
    if data.len() < 4 + secret_len {
        return Err(MlsEngineError::Client("signing key file truncated".into()));
    }
    let secret_bytes = Zeroizing::new(data[4..4 + secret_len].to_vec());
    let public_bytes = data[4 + secret_len..].to_vec();
    Ok((
        SignatureSecretKey::from(secret_bytes.as_slice().to_vec()),
        SignaturePublicKey::from(public_bytes),
    ))
}

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

    /// Creating an MLS engine with a DID URI identity produces a CBOR-encoded
    /// credential (version 1) rather than raw bytes.
    #[test]
    fn test_did_identity_produces_cbor_credential() {
        let dir = tempfile::tempdir().expect("failed to create temp dir");

        let did_uri = b"did:web:agent.example.com";
        let engine = MlsEngine::new(dir.path(), did_uri).expect("engine creation failed");

        // Generate a key package so we can inspect the credential inside it.
        let kp_bytes = engine.generate_key_package().expect("key package failed");

        // The key package embeds the BasicCredential which holds our identity
        // bytes. Parse the credential from the key package.
        let kp_msg = MlsMessage::from_bytes(&kp_bytes).expect("parse key package");
        let kp = kp_msg.into_key_package().expect("extract key package");
        let basic = kp
            .signing_identity()
            .credential
            .as_basic()
            .expect("expected BasicCredential");

        // The credential identity should be CBOR-encoded, not raw UTF-8.
        let parsed = SkytaleIdentity::parse(&basic.identifier);
        match parsed {
            SkytaleIdentity::DidWeb { uri, pubkey } => {
                assert_eq!(uri, "did:web:agent.example.com");
                // The pubkey should be 32 bytes (Ed25519).
                assert_eq!(pubkey.len(), 32);
                // Pubkey should not be all zeros (it's a real generated key).
                assert!(pubkey.iter().any(|&b| b != 0));
            }
            other => {
                panic!("expected DidWeb identity, got {other:?}");
            }
        }
    }

    /// Creating an MLS engine with plain identity bytes still works as RawKey.
    #[test]
    fn test_raw_identity_remains_raw() {
        let dir = tempfile::tempdir().expect("failed to create temp dir");

        let raw_id = b"agent-plain";
        let engine = MlsEngine::new(dir.path(), raw_id).expect("engine creation failed");

        let kp_bytes = engine.generate_key_package().expect("key package failed");
        let kp_msg = MlsMessage::from_bytes(&kp_bytes).expect("parse key package");
        let kp = kp_msg.into_key_package().expect("extract key package");
        let basic = kp
            .signing_identity()
            .credential
            .as_basic()
            .expect("expected BasicCredential");

        let parsed = SkytaleIdentity::parse(&basic.identifier);
        match parsed {
            SkytaleIdentity::RawKey(bytes) => {
                assert_eq!(bytes, b"agent-plain");
            }
            other => {
                panic!("expected RawKey identity, got {other:?}");
            }
        }
    }

    /// Round-trip: DidWeb identity serializes to CBOR and parses back correctly.
    #[test]
    fn test_did_identity_roundtrip() {
        let pubkey = [42u8; 32];
        let identity = SkytaleIdentity::DidWeb {
            uri: "did:web:test.example.org".to_string(),
            pubkey,
        };

        let bytes = identity.to_credential_bytes();
        let parsed = SkytaleIdentity::parse(&bytes);
        assert_eq!(parsed, identity);
    }

    /// Creating an MLS engine with a DID:key URI identity produces a CBOR-encoded
    /// credential (version 2) with the DidKey variant.
    #[test]
    fn test_did_key_identity_produces_cbor_credential() {
        let dir = tempfile::tempdir().expect("failed to create temp dir");

        let did_uri = b"did:key:z6MkhaXgBZDvotDkL5257faiztiGiC2QtKLGpbnnEGta2doK";
        let engine = MlsEngine::new(dir.path(), did_uri).expect("engine creation failed");

        let kp_bytes = engine.generate_key_package().expect("key package failed");
        let kp_msg = MlsMessage::from_bytes(&kp_bytes).expect("parse key package");
        let kp = kp_msg.into_key_package().expect("extract key package");
        let basic = kp
            .signing_identity()
            .credential
            .as_basic()
            .expect("expected BasicCredential");

        let parsed = SkytaleIdentity::parse(&basic.identifier);
        match parsed {
            SkytaleIdentity::DidKey { did_uri, pubkey } => {
                assert_eq!(
                    did_uri,
                    "did:key:z6MkhaXgBZDvotDkL5257faiztiGiC2QtKLGpbnnEGta2doK"
                );
                assert_eq!(pubkey.len(), 32);
                // Pubkey should not be all zeros (it's a real generated key).
                assert!(pubkey.iter().any(|&b| b != 0));
            }
            other => {
                panic!("expected DidKey identity, got {other:?}");
            }
        }
    }

    /// Round-trip: DidKey identity serializes to CBOR and parses back correctly.
    #[test]
    fn test_did_key_identity_roundtrip() {
        let pubkey = [77u8; 32];
        let identity = SkytaleIdentity::DidKey {
            did_uri: "did:key:z6MkhaXgBZDvotDkL5257faiztiGiC2QtKLGpbnnEGta2doK".to_string(),
            pubkey,
        };

        let bytes = identity.to_credential_bytes();
        let parsed = SkytaleIdentity::parse(&bytes);
        assert_eq!(parsed, identity);
    }

    /// rotate_key produces a serialized commit message and leaves the commit pending.
    #[test]
    fn test_rotate_key_produces_commit() {
        let dir = tempfile::tempdir().expect("failed to create temp dir");
        let mut engine = MlsEngine::new(dir.path(), b"rotate-test").expect("engine failed");

        let group_id = engine
            .create_group(b"rotate-group")
            .expect("create group failed");
        let commit_bytes = engine.rotate_key(&group_id).expect("rotate_key failed");

        // The commit bytes must be a valid MLS message.
        MlsMessage::from_bytes(&commit_bytes).expect("commit should be valid MLS message");

        // Apply the pending commit to advance the epoch.
        engine
            .apply_pending_commit(&group_id)
            .expect("apply pending commit failed");
    }

    /// export_audit_key returns a 32-byte secret derived from the MLS exporter.
    #[test]
    fn test_export_audit_key_returns_32_bytes() {
        let dir = tempfile::tempdir().expect("failed to create temp dir");
        let mut engine = MlsEngine::new(dir.path(), b"audit-test").expect("engine failed");

        let group_id = engine
            .create_group(b"audit-group")
            .expect("create group failed");
        let audit_key = engine.export_audit_key(&group_id).expect("export failed");

        assert_eq!(audit_key.len(), 32);
        // The key should not be all zeros (statistically impossible for a real derivation).
        assert!(audit_key.iter().any(|&b| b != 0));
    }

    /// export_audit_key produces different keys for different groups.
    #[test]
    fn test_export_audit_key_differs_across_groups() {
        let dir = tempfile::tempdir().expect("failed to create temp dir");
        let mut engine = MlsEngine::new(dir.path(), b"audit-multi").expect("engine failed");

        let gid1 = engine
            .create_group(b"audit-group-1")
            .expect("create group 1");
        let gid2 = engine
            .create_group(b"audit-group-2")
            .expect("create group 2");

        let key1 = engine.export_audit_key(&gid1).expect("export 1");
        let key2 = engine.export_audit_key(&gid2).expect("export 2");

        assert_ne!(key1, key2, "audit keys for different groups must differ");
    }
}
