//! KeyManager implementation for MEGA ^!keys attribute (simplified but format-compatible).
//!
//! This mirrors the SDK's `KeyManager::toKeysContainer/fromKeysContainer` layout enough to keep
//! share/export working on upgraded accounts. It supports the core fields (version/generation,
//! identity, privkeys, authrings, share keys, pending shares) and encrypts the container with
//! AES-128-GCM using the master key.

use std::collections::{HashMap, HashSet};

use aes_gcm::aead::{Aead, KeyInit, OsRng};
use aes_gcm::{Aes128Gcm, Nonce};
use hmac::{Hmac, Mac};
use rand::RngCore;
use sha2::Sha256;

use crate::base64::base64url_decode;
use crate::crypto::AuthRing;
use crate::error::{MegaError, Result};

// Tag constants (from SDK)
const TAG_VERSION: u8 = 1;
const TAG_CREATION_TIME: u8 = 2;
const TAG_IDENTITY: u8 = 3;
const TAG_GENERATION: u8 = 4;
const TAG_ATTR: u8 = 5;
const TAG_PRIV_ED25519: u8 = 16;
const TAG_PRIV_CU25519: u8 = 17;
const TAG_PRIV_RSA: u8 = 18;
const TAG_AUTHRING_ED25519: u8 = 32;
const TAG_AUTHRING_CU25519: u8 = 33;
const TAG_SHAREKEYS: u8 = 48;
const TAG_PENDING_OUTSHARES: u8 = 64;
const TAG_PENDING_INSHARES: u8 = 65;
const TAG_BACKUPS: u8 = 80;
const TAG_WARNINGS: u8 = 96;

const HEADER_BYTE0: u8 = 20;
const HEADER_BYTE1: u8 = 0;
const GCM_IV_LEN: usize = 12;
/// Share key flag indicating a trusted key.
pub const SHAREKEY_FLAG_TRUSTED: u8 = 1 << 0;
/// Share key flag indicating the key is actively in use.
pub const SHAREKEY_FLAG_IN_USE: u8 = 1 << 1;

/// Share key entry for a node handle.
///
/// `handle` is the 6-byte node handle, `key` is the 16-byte share key, and
/// `flags` stores `SHAREKEY_FLAG_*` bits.
#[derive(Debug, Clone, Default)]
pub struct ShareKeyEntry {
    pub handle: [u8; 6],
    pub key: [u8; 16],
    pub flags: u8,
}

/// Pending outgoing share entry.
///
/// Stores the target folder handle and recipient identifier.
#[derive(Debug, Clone)]
pub struct PendingOutEntry {
    pub node_handle: [u8; 6],
    pub uid: PendingUid,
}

#[derive(Debug, Clone)]
/// Identifier for a pending share recipient.
pub enum PendingUid {
    /// Pending share for a user handle (8 bytes).
    UserHandle([u8; 8]),
    /// Pending share for an email address.
    Email(String),
}

/// Pending incoming share entry.
///
/// Maps a node handle to the sender handle and the encrypted share key payload.
#[derive(Debug, Clone)]
pub struct PendingInEntry {
    pub node_handle_b64: String,
    pub user_handle: [u8; 8],
    pub share_key: Vec<u8>,
}

/// Warnings map carried inside ^!keys.
///
/// Entries are stored as `(tag, value)` pairs. Helpers focus on the `cv` tag
/// (contact verification requirement).
///
/// # Examples
/// ```
/// use megalib::crypto::Warnings;
///
/// let mut warnings = Warnings::default();
/// warnings.set_cv(true);
/// assert!(warnings.cv_enabled());
/// ```
#[derive(Debug, Clone, Default)]
pub struct Warnings(pub Vec<(String, Vec<u8>)>);

impl Warnings {
    /// Set the contact verification warning (`cv`) flag.
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::Warnings;
    ///
    /// let mut warnings = Warnings::default();
    /// warnings.set_cv(true);
    /// assert!(warnings.cv_enabled());
    /// ```
    pub fn set_cv(&mut self, enabled: bool) {
        let val = if enabled {
            b"1".to_vec()
        } else {
            b"0".to_vec()
        };
        if let Some(entry) = self.0.iter_mut().find(|(k, _)| k == "cv") {
            entry.1 = val;
        } else {
            self.0.push(("cv".to_string(), val));
        }
    }

    /// Check whether the contact verification warning is enabled.
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::Warnings;
    ///
    /// let mut warnings = Warnings::default();
    /// warnings.set_cv(false);
    /// assert!(!warnings.cv_enabled());
    /// ```
    pub fn cv_enabled(&self) -> bool {
        self.0
            .iter()
            .find(|(k, _)| k == "cv")
            .map(|(_, v)| v != b"0")
            .unwrap_or(false)
    }
}

/// In-memory representation of the ^!keys container.
///
/// This struct stores key material, authrings, share keys, and related metadata
/// needed to serialize or update MEGA's ^!keys attribute.
#[derive(Debug, Clone, Default)]
pub struct KeyManager {
    pub version: u8,
    pub creation_time: u32,
    pub identity: u64,
    pub generation: u32,
    pub attr: Vec<u8>,
    pub priv_ed25519: Vec<u8>,
    pub priv_cu25519: Vec<u8>,
    pub priv_rsa: Vec<u8>,
    pub authring_ed: AuthRing,
    pub authring_cu: AuthRing,
    pub share_keys: Vec<ShareKeyEntry>,
    /// O(1) handle→index into `share_keys` for fast lookups.
    share_key_index: HashMap<[u8; 6], usize>,
    pub pending_out: Vec<PendingOutEntry>,
    pub pending_in: Vec<PendingInEntry>,
    pub backups: Vec<u8>,
    pub warnings: Warnings,

    /// When true, the user must manually verify contacts before share-key exchange (mirrors SDK flag).
    pub manual_verification: bool,
}

/// Derive the ^!keys AES-128-GCM key from the master key using HKDF-SHA256 (info byte = 1).
fn derive_keys_cipher(master_key: &[u8; 16]) -> [u8; 16] {
    // HKDF-Extract with salt = zeros(hashlen)
    let mut prk_mac = <Hmac<Sha256> as Mac>::new_from_slice(&[0u8; 32]).expect("HMAC key len");
    prk_mac.update(master_key);
    let prk = prk_mac.finalize().into_bytes();

    // HKDF-Expand with info = {1}, single block needed for 16 bytes
    let mut okm_mac = <Hmac<Sha256> as Mac>::new_from_slice(&prk).expect("HMAC key len");
    okm_mac.update(&[1u8]); // info
    okm_mac.update(&[1u8]); // counter
    let t1 = okm_mac.finalize().into_bytes();

    let mut out = [0u8; 16];
    out.copy_from_slice(&t1[..16]);
    out
}

impl KeyManager {
    fn validate_authring(blob: &[u8]) -> Result<()> {
        if blob.len() > 64 * 1024 {
            return Err(MegaError::InvalidResponse);
        }
        // Must be decodable as LTLV map
        let _ = Self::parse_ltlv_map(blob)?;
        Ok(())
    }

    /// Rebuild the handle→index HashMap from the `share_keys` Vec.
    fn rebuild_share_key_index(&mut self) {
        self.share_key_index.clear();
        for (i, sk) in self.share_keys.iter().enumerate() {
            self.share_key_index.insert(sk.handle, i);
        }
    }

    /// Create a new key manager with default metadata.
    ///
    /// This sets `version = 1`, `generation = 0`, and clears the manual
    /// verification flag.
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::KeyManager;
    ///
    /// let km = KeyManager::new();
    /// assert!(!km.is_ready());
    /// ```
    pub fn new() -> Self {
        KeyManager {
            version: 1,
            generation: 0,
            manual_verification: false,
            ..KeyManager::default()
        }
    }

    /// Check whether private Ed25519 and Cu25519 keys are present.
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// assert!(!km.is_ready());
    /// km.set_priv_keys(&[1u8; 32], &[2u8; 32]);
    /// assert!(km.is_ready());
    /// ```
    pub fn is_ready(&self) -> bool {
        !self.priv_ed25519.is_empty() && !self.priv_cu25519.is_empty()
    }

    /// Store Ed25519 and Cu25519 private keys.
    ///
    /// The input slices must be 32 bytes each.
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// km.set_priv_keys(&[1u8; 32], &[2u8; 32]);
    /// assert!(km.is_ready());
    /// ```
    pub fn set_priv_keys(&mut self, ed: &[u8], cu: &[u8]) {
        self.priv_ed25519 = ed.to_vec();
        self.priv_cu25519 = cu.to_vec();
    }

    /// Insert or update a share key with explicit flag bits.
    ///
    /// Invalid handles or keys are ignored.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// let handle = base64url_encode(&[0u8; 6]);
    /// km.add_share_key_with_flags(&handle, &[7u8; 16], true, false);
    /// assert!(km.is_share_key_trusted(&handle));
    /// ```
    pub fn add_share_key_with_flags(
        &mut self,
        handle_b64: &str,
        key: &[u8],
        trusted: bool,
        in_use: bool,
    ) {
        if key.len() != 16 {
            return;
        }
        let Some(handle) = Self::decode_handle(handle_b64) else {
            return;
        };
        let mut flags = 0u8;
        if trusted {
            flags |= SHAREKEY_FLAG_TRUSTED;
        }
        if in_use {
            flags |= SHAREKEY_FLAG_IN_USE;
        }

        if let Some(&idx) = self.share_key_index.get(&handle)
            && let Some(entry) = self.share_keys.get_mut(idx)
        {
            entry.key.copy_from_slice(key);
            entry.flags |= flags;
            return;
        }

        let mut k = [0u8; 16];
        k.copy_from_slice(key);
        let idx = self.share_keys.len();
        self.share_keys.push(ShareKeyEntry {
            handle,
            key: k,
            flags,
        });
        self.share_key_index.insert(handle, idx);
    }

    /// Insert or update a share key using a base64url handle.
    ///
    /// The handle must decode to 6 bytes and the key must be 16 bytes.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// let handle = base64url_encode(&[0u8; 6]);
    /// km.add_share_key_from_str(&handle, &[1u8; 16]);
    /// assert!(km.get_share_key_from_str(&handle).is_some());
    /// ```
    pub fn add_share_key_from_str(&mut self, handle_b64: &str, key: &[u8]) {
        self.add_share_key_with_flags(handle_b64, key, false, false);
    }

    /// Insert or update a share key using owned handle and key data.
    ///
    /// Invalid handles or keys are ignored.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// let handle = base64url_encode(&[2u8; 6]);
    /// km.add_share_key(handle.clone(), vec![9u8; 16]);
    /// assert!(km.get_share_key_from_str(&handle).is_some());
    /// ```
    pub fn add_share_key(&mut self, handle_b64: String, key: Vec<u8>) {
        self.add_share_key_with_flags(&handle_b64, &key, false, false);
    }

    /// Return the raw flag bits for a share key.
    ///
    /// Returns `None` if the handle is invalid or not present.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// let handle = base64url_encode(&[3u8; 6]);
    /// km.add_share_key_from_str(&handle, &[1u8; 16]);
    /// assert!(km.share_key_flags(&handle).is_some());
    /// ```
    pub fn share_key_flags(&self, handle_b64: &str) -> Option<u8> {
        let decoded = Self::decode_handle(handle_b64)?;
        self.share_key_index
            .get(&decoded)
            .and_then(|&idx| self.share_keys.get(idx))
            .map(|e| e.flags)
    }

    /// Check whether a share key is marked trusted.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// let handle = base64url_encode(&[4u8; 6]);
    /// km.add_share_key_with_flags(&handle, &[2u8; 16], true, false);
    /// assert!(km.is_share_key_trusted(&handle));
    /// ```
    pub fn is_share_key_trusted(&self, handle_b64: &str) -> bool {
        self.share_key_flags(handle_b64)
            .map(|f| f & SHAREKEY_FLAG_TRUSTED != 0)
            .unwrap_or(false)
    }

    /// Check whether a share key is marked in use.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// let handle = base64url_encode(&[5u8; 6]);
    /// km.add_share_key_with_flags(&handle, &[3u8; 16], false, true);
    /// assert!(km.is_share_key_in_use(&handle));
    /// ```
    pub fn is_share_key_in_use(&self, handle_b64: &str) -> bool {
        self.share_key_flags(handle_b64)
            .map(|f| f & SHAREKEY_FLAG_IN_USE != 0)
            .unwrap_or(false)
    }

    /// Set or clear the IN_USE flag for a share key.
    ///
    /// Returns `false` if the handle is not present.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// let handle = base64url_encode(&[6u8; 6]);
    /// km.add_share_key_from_str(&handle, &[4u8; 16]);
    /// assert!(km.set_share_key_in_use(&handle, true));
    /// assert!(km.is_share_key_in_use(&handle));
    /// ```
    pub fn set_share_key_in_use(&mut self, handle_b64: &str, in_use: bool) -> bool {
        self.set_share_key_flag(handle_b64, SHAREKEY_FLAG_IN_USE, in_use)
    }

    /// Set or clear the TRUSTED flag for a share key.
    ///
    /// Returns `false` if the handle is not present.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// let handle = base64url_encode(&[7u8; 6]);
    /// km.add_share_key_from_str(&handle, &[5u8; 16]);
    /// assert!(km.set_share_key_trusted(&handle, true));
    /// assert!(km.is_share_key_trusted(&handle));
    /// ```
    pub fn set_share_key_trusted(&mut self, handle_b64: &str, trusted: bool) -> bool {
        self.set_share_key_flag(handle_b64, SHAREKEY_FLAG_TRUSTED, trusted)
    }

    /// Enable or disable manual verification gating.
    ///
    /// When enabled, contacts must be manually verified before share-key exchange.
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// km.set_manual_verification(true);
    /// assert!(km.manual_verification());
    /// ```
    pub fn set_manual_verification(&mut self, enabled: bool) {
        self.manual_verification = enabled;
    }

    /// Check whether manual verification gating is enabled.
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::KeyManager;
    ///
    /// let km = KeyManager::new();
    /// assert!(!km.manual_verification());
    /// ```
    pub fn manual_verification(&self) -> bool {
        self.manual_verification
    }

    /// Update the contact verification warning flag (`cv`).
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// km.set_contact_verification_warning(true);
    /// assert!(km.contact_verification_warning());
    /// ```
    pub fn set_contact_verification_warning(&mut self, enabled: bool) {
        let val = if enabled {
            b"1".to_vec()
        } else {
            b"0".to_vec()
        };
        if let Some(entry) = self.warnings.0.iter_mut().find(|(k, _)| k == "cv") {
            entry.1 = val;
        } else {
            self.warnings.0.push(("cv".to_string(), val));
        }
    }

    /// Check whether the contact verification warning flag is set.
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// km.set_contact_verification_warning(false);
    /// assert!(!km.contact_verification_warning());
    /// ```
    pub fn contact_verification_warning(&self) -> bool {
        self.warnings
            .0
            .iter()
            .find(|(k, _)| k == "cv")
            .map(|(_, v)| v != b"0")
            .unwrap_or(false)
    }

    /// Set the raw backups blob.
    ///
    /// The blob is stored verbatim and emitted into the `^!keys` container
    /// when calling [`KeyManager::encode_container`].
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// km.set_backups(vec![1, 2, 3]);
    /// ```
    pub fn set_backups(&mut self, blob: Vec<u8>) {
        self.backups = blob;
    }

    /// Replace the Ed25519 authring from an LTLV blob.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::{AuthRing, KeyManager};
    ///
    /// let handle = base64url_encode(&[1u8; 8]);
    /// let mut ring = AuthRing::default();
    /// ring.update(&handle, b"public-key", true);
    ///
    /// let mut km = KeyManager::new();
    /// km.set_authring_ed25519(ring.serialize_ltlv());
    /// ```
    pub fn set_authring_ed25519(&mut self, blob: Vec<u8>) {
        self.authring_ed = AuthRing::deserialize_ltlv(&blob);
    }

    /// Replace the Cu25519 authring from an LTLV blob.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::{AuthRing, KeyManager};
    ///
    /// let handle = base64url_encode(&[2u8; 8]);
    /// let mut ring = AuthRing::default();
    /// ring.update(&handle, b"cu25519-key", false);
    ///
    /// let mut km = KeyManager::new();
    /// km.set_authring_cu25519(ring.serialize_ltlv());
    /// ```
    pub fn set_authring_cu25519(&mut self, blob: Vec<u8>) {
        self.authring_cu = AuthRing::deserialize_ltlv(&blob);
    }

    /// Borrow the warnings map.
    ///
    /// The warnings map is serialized into the `^!keys` container. Use
    /// [`Warnings::set_cv`] to toggle the contact verification flag.
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::KeyManager;
    ///
    /// let km = KeyManager::new();
    /// assert!(!km.warnings().cv_enabled());
    /// ```
    pub fn warnings(&self) -> &Warnings {
        &self.warnings
    }

    /// Replace the warnings map.
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::{KeyManager, Warnings};
    ///
    /// let mut km = KeyManager::new();
    /// km.set_warnings(Warnings::default());
    /// ```
    pub fn set_warnings(&mut self, warnings: Warnings) {
        self.warnings = warnings;
    }

    /// Get a share key for a handle encoded as base64url.
    ///
    /// Returns `None` if the handle is invalid or not present.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// let handle = base64url_encode(&[8u8; 6]);
    /// km.add_share_key_from_str(&handle, &[6u8; 16]);
    /// assert!(km.get_share_key_from_str(&handle).is_some());
    /// ```
    /// Check if a share key is present for the given base64url handle.
    pub fn contains_share_key(&self, handle_b64: &str) -> bool {
        self.get_share_key_from_str(handle_b64).is_some()
    }

    pub fn get_share_key_from_str(&self, handle_b64: &str) -> Option<[u8; 16]> {
        let decoded = base64url_decode(handle_b64).ok()?;
        if decoded.len() != 6 {
            return None;
        }
        let mut h = [0u8; 6];
        h.copy_from_slice(&decoded);
        self.share_key_index
            .get(&h)
            .and_then(|&idx| self.share_keys.get(idx))
            .map(|sk| sk.key)
    }

    /// Add a pending outgoing share for an email recipient.
    ///
    /// Invalid emails or duplicate entries are ignored.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// let handle = base64url_encode(&[9u8; 6]);
    /// km.add_pending_out_email(&handle, "user@example.com");
    /// ```
    pub fn add_pending_out_email(&mut self, handle_b64: &str, email: &str) {
        if email.is_empty() || email.len() >= 256 {
            return;
        }
        let Some(handle) = Self::decode_handle(handle_b64) else {
            return;
        };
        let uid = PendingUid::Email(email.to_string());
        if self
            .pending_out
            .iter()
            .any(|e| e.node_handle == handle && matches_pending_uid(&e.uid, &uid))
        {
            return;
        }
        self.pending_out.push(PendingOutEntry {
            node_handle: handle,
            uid,
        });
    }

    /// Add a pending outgoing share for a user handle.
    ///
    /// Duplicate entries are ignored.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// let handle = base64url_encode(&[10u8; 6]);
    /// km.add_pending_out_user_handle(&handle, &[1u8; 8]);
    /// ```
    pub fn add_pending_out_user_handle(&mut self, handle_b64: &str, user_handle: &[u8; 8]) {
        let Some(handle) = Self::decode_handle(handle_b64) else {
            return;
        };
        let uid = PendingUid::UserHandle(*user_handle);
        if self
            .pending_out
            .iter()
            .any(|e| e.node_handle == handle && matches_pending_uid(&e.uid, &uid))
        {
            return;
        }
        self.pending_out.push(PendingOutEntry {
            node_handle: handle,
            uid,
        });
    }

    /// Remove a pending outgoing share entry.
    ///
    /// Returns `true` if an entry was removed.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::{KeyManager, PendingUid};
    ///
    /// let mut km = KeyManager::new();
    /// let handle = base64url_encode(&[11u8; 6]);
    /// km.add_pending_out_email(&handle, "user@example.com");
    /// let removed = km.remove_pending_out(&handle, &PendingUid::Email("user@example.com".into()));
    /// assert!(removed);
    /// ```
    pub fn remove_pending_out(&mut self, handle_b64: &str, uid: &PendingUid) -> bool {
        let Some(handle) = Self::decode_handle(handle_b64) else {
            return false;
        };
        let before = self.pending_out.len();
        self.pending_out
            .retain(|e| !(e.node_handle == handle && matches_pending_uid(&e.uid, uid)));
        before != self.pending_out.len()
    }

    /// Add a pending incoming share entry.
    ///
    /// Empty share keys and duplicates are ignored.
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// km.add_pending_in("HANDLE", &[2u8; 8], vec![0u8; 16]);
    /// ```
    pub fn add_pending_in(&mut self, handle_b64: &str, user_handle: &[u8; 8], share_key: Vec<u8>) {
        if share_key.is_empty() {
            return;
        }
        if self
            .pending_in
            .iter()
            .any(|e| e.node_handle_b64 == handle_b64 && e.user_handle == *user_handle)
        {
            return;
        }
        self.pending_in.push(PendingInEntry {
            node_handle_b64: handle_b64.to_string(),
            user_handle: *user_handle,
            share_key,
        });
    }

    /// Remove a pending incoming share entry.
    ///
    /// Returns `true` if an entry was removed.
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// km.add_pending_in("HANDLE", &[3u8; 8], vec![1u8; 16]);
    /// assert!(km.remove_pending_in("HANDLE", &[3u8; 8]));
    /// ```
    pub fn remove_pending_in(&mut self, handle_b64: &str, user_handle: &[u8; 8]) -> bool {
        let before = self.pending_in.len();
        self.pending_in
            .retain(|e| !(e.node_handle_b64 == handle_b64 && e.user_handle == *user_handle));
        before != self.pending_in.len()
    }

    /// Clear the IN_USE flag for share keys not present in the provided handle set.
    ///
    /// Returns `true` if any flags were changed.
    ///
    /// # Examples
    /// ```
    /// use std::collections::HashSet;
    ///
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// let handle = base64url_encode(&[12u8; 6]);
    /// km.add_share_key_with_flags(&handle, &[1u8; 16], false, true);
    ///
    /// let keep = HashSet::new();
    /// assert!(km.clear_in_use_not_in(&keep));
    /// ```
    pub fn clear_in_use_not_in(&mut self, handles_b64: &HashSet<String>) -> bool {
        let mut changed = false;
        for sk in &mut self.share_keys {
            let handle_b64 = crate::base64::base64url_encode(&sk.handle);
            if !handles_b64.contains(&handle_b64) && (sk.flags & SHAREKEY_FLAG_IN_USE != 0) {
                sk.flags &= !SHAREKEY_FLAG_IN_USE;
                changed = true;
            }
        }
        changed
    }

    /// Clear the TRUSTED flag for share keys not present in the provided handle set.
    ///
    /// Returns `true` if any flags were changed.
    ///
    /// # Examples
    /// ```
    /// use std::collections::HashSet;
    ///
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// let handle = base64url_encode(&[13u8; 6]);
    /// km.add_share_key_with_flags(&handle, &[2u8; 16], true, false);
    ///
    /// let keep = HashSet::new();
    /// assert!(km.clear_trusted_not_in(&keep));
    /// ```
    pub fn clear_trusted_not_in(&mut self, handles_b64: &HashSet<String>) -> bool {
        let mut changed = false;
        for sk in &mut self.share_keys {
            let handle_b64 = crate::base64::base64url_encode(&sk.handle);
            if !handles_b64.contains(&handle_b64) && (sk.flags & SHAREKEY_FLAG_TRUSTED != 0) {
                sk.flags &= !SHAREKEY_FLAG_TRUSTED;
                changed = true;
            }
        }
        changed
    }

    fn set_share_key_flag(&mut self, handle_b64: &str, flag: u8, enabled: bool) -> bool {
        let Some(handle) = Self::decode_handle(handle_b64) else {
            return false;
        };
        if let Some(&idx) = self.share_key_index.get(&handle)
            && let Some(entry) = self.share_keys.get_mut(idx)
        {
            if enabled {
                entry.flags |= flag;
            } else {
                entry.flags &= !flag;
            }
            return true;
        }
        false
    }

    fn decode_handle(handle_b64: &str) -> Option<[u8; 6]> {
        let decoded = base64url_decode(handle_b64).ok()?;
        if decoded.len() != 6 {
            return None;
        }
        let mut h = [0u8; 6];
        h.copy_from_slice(&decoded);
        Some(h)
    }

    /// Remove a share key entry (and associated flags) by handle.
    ///
    /// Returns `false` if the handle is invalid or not present.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut km = KeyManager::new();
    /// let handle = base64url_encode(&[14u8; 6]);
    /// km.add_share_key_from_str(&handle, &[1u8; 16]);
    /// assert!(km.remove_share_key(&handle));
    /// ```
    pub fn remove_share_key(&mut self, handle_b64: &str) -> bool {
        if let Some(handle) = Self::decode_handle(handle_b64) {
            let before = self.share_keys.len();
            self.share_keys.retain(|e| e.handle != handle);
            if self.share_keys.len() != before {
                self.rebuild_share_key_index();
                return true;
            }
        }
        false
    }

    /// Clear TRUSTED flag on all share keys.
    pub fn clear_trusted_all(&mut self) -> bool {
        let mut changed = false;
        for sk in &mut self.share_keys {
            if (sk.flags & SHAREKEY_FLAG_TRUSTED) != 0 {
                sk.flags &= !SHAREKEY_FLAG_TRUSTED;
                changed = true;
            }
        }
        changed
    }
    /// Encode the key manager into an encrypted `^!keys` container.
    ///
    /// The state is serialized to LTLV and encrypted with AES-128-GCM using a
    /// key derived from `master_key`.
    ///
    /// # Errors
    ///
    /// Returns an error if the stored RSA private key is shorter than 512 bytes,
    /// or if encryption fails.
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::KeyManager;
    ///
    /// let master = [0u8; 16];
    /// let mut km = KeyManager::new();
    /// km.set_priv_keys(&[1u8; 32], &[2u8; 32]);
    ///
    /// let blob = km.encode_container(&master).expect("encode");
    /// assert!(blob.len() > 2);
    /// ```
    pub fn encode_container(&self, master_key: &[u8; 16]) -> Result<Vec<u8>> {
        if !self.priv_rsa.is_empty() && self.priv_rsa.len() < 512 {
            return Err(MegaError::Custom(
                "Invalid RSA key length in ^!keys (expected empty or >=512 bytes)".into(),
            ));
        }
        let plain = self.serialize_ltlv()?;
        // GCM encrypt
        let derived = derive_keys_cipher(master_key);
        let cipher = Aes128Gcm::new_from_slice(&derived)
            .map_err(|e| MegaError::CryptoError(format!("GCM init: {}", e)))?;
        let mut iv = [0u8; GCM_IV_LEN];
        OsRng.fill_bytes(&mut iv);
        let nonce = Nonce::from_slice(&iv);
        let mut ct = cipher
            .encrypt(nonce, plain.as_ref())
            .map_err(|e| MegaError::CryptoError(format!("GCM encrypt: {}", e)))?;

        let mut out = Vec::with_capacity(2 + GCM_IV_LEN + ct.len());
        out.push(HEADER_BYTE0);
        out.push(HEADER_BYTE1);
        out.extend_from_slice(&iv);
        out.append(&mut ct);
        Ok(out)
    }

    /// Decode and apply an encrypted `^!keys` container.
    ///
    /// The blob is validated, decrypted with AES-128-GCM, and parsed as LTLV.
    /// On success the current state is replaced. If the incoming generation is
    /// older than the current generation, the decode is rejected.
    ///
    /// # Errors
    ///
    /// Returns an error if the blob is malformed, decryption fails, the LTLV
    /// payload is invalid, or a generation downgrade is detected.
    ///
    /// # Examples
    /// ```
    /// use megalib::crypto::KeyManager;
    ///
    /// let master = [0u8; 16];
    /// let mut km = KeyManager::new();
    /// km.set_priv_keys(&[1u8; 32], &[2u8; 32]);
    /// let blob = km.encode_container(&master).expect("encode");
    ///
    /// let mut decoded = KeyManager::new();
    /// decoded.decode_container(&blob, &master).expect("decode");
    /// assert!(decoded.is_ready());
    /// ```
    pub fn decode_container(&mut self, data: &[u8], master_key: &[u8; 16]) -> Result<()> {
        if data.len() < 2 + GCM_IV_LEN + 16 {
            return Err(MegaError::InvalidResponse);
        }
        if data[0] != HEADER_BYTE0 || data[1] != HEADER_BYTE1 {
            return Err(MegaError::InvalidResponse);
        }
        let iv = &data[2..2 + GCM_IV_LEN];
        let ct = &data[2 + GCM_IV_LEN..];
        let derived = derive_keys_cipher(master_key);
        let cipher = Aes128Gcm::new_from_slice(&derived)
            .map_err(|e| MegaError::CryptoError(format!("GCM init: {}", e)))?;
        let nonce = Nonce::from_slice(iv);
        let plain = cipher
            .decrypt(nonce, ct)
            .map_err(|e| MegaError::CryptoError(format!("GCM decrypt: {}", e)))?;
        // Parse payload
        let mut tmp = self.clone();
        tmp.deserialize_ltlv(&plain)?;

        // Downgrade protection: reject if received generation lower than current
        if tmp.generation < self.generation {
            return Err(MegaError::DowngradeDetected);
        }

        *self = tmp;
        Ok(())
    }

    fn tag_header(tag: u8, len: usize) -> [u8; 4] {
        [
            tag,
            ((len >> 16) & 0xFF) as u8,
            ((len >> 8) & 0xFF) as u8,
            (len & 0xFF) as u8,
        ]
    }

    fn serialize_ltlv(&self) -> Result<Vec<u8>> {
        let mut out = Vec::new();
        // version
        out.extend_from_slice(&Self::tag_header(TAG_VERSION, 1));
        out.push(self.version);
        // creation time
        let ct_be = self.creation_time.to_be_bytes();
        out.extend_from_slice(&Self::tag_header(TAG_CREATION_TIME, ct_be.len()));
        out.extend_from_slice(&ct_be);
        // identity
        let id_bytes = self.identity.to_le_bytes(); // SDK stores as little? It appends raw u64.
        out.extend_from_slice(&Self::tag_header(TAG_IDENTITY, id_bytes.len()));
        out.extend_from_slice(&id_bytes);
        // generation stored on wire as (gen + 1) like SDK
        let gen_be = (self.generation + 1).to_be_bytes();
        out.extend_from_slice(&Self::tag_header(TAG_GENERATION, gen_be.len()));
        out.extend_from_slice(&gen_be);
        // attr
        out.extend_from_slice(&Self::tag_header(TAG_ATTR, self.attr.len()));
        out.extend_from_slice(&self.attr);
        // priv ed
        out.extend_from_slice(&Self::tag_header(TAG_PRIV_ED25519, self.priv_ed25519.len()));
        out.extend_from_slice(&self.priv_ed25519);
        // priv cu
        out.extend_from_slice(&Self::tag_header(TAG_PRIV_CU25519, self.priv_cu25519.len()));
        out.extend_from_slice(&self.priv_cu25519);
        // priv rsa
        out.extend_from_slice(&Self::tag_header(TAG_PRIV_RSA, self.priv_rsa.len()));
        out.extend_from_slice(&self.priv_rsa);
        // authrings (serialize AuthRing objects to LTLV blobs)
        let auth_ed_blob = self.authring_ed.serialize_ltlv();
        out.extend_from_slice(&Self::tag_header(TAG_AUTHRING_ED25519, auth_ed_blob.len()));
        out.extend_from_slice(&auth_ed_blob);
        let auth_cu_blob = self.authring_cu.serialize_ltlv();
        out.extend_from_slice(&Self::tag_header(TAG_AUTHRING_CU25519, auth_cu_blob.len()));
        out.extend_from_slice(&auth_cu_blob);
        // share keys
        let sk_blob = self.serialize_share_keys();
        out.extend_from_slice(&Self::tag_header(TAG_SHAREKEYS, sk_blob.len()));
        out.extend_from_slice(&sk_blob);
        // pending out/in, backups, warnings (may be empty)
        let po_blob = self.serialize_pending_out();
        out.extend_from_slice(&Self::tag_header(TAG_PENDING_OUTSHARES, po_blob.len()));
        out.extend_from_slice(&po_blob);
        let pi_blob = self.serialize_pending_in();
        out.extend_from_slice(&Self::tag_header(TAG_PENDING_INSHARES, pi_blob.len()));
        out.extend_from_slice(&pi_blob);
        out.extend_from_slice(&Self::tag_header(TAG_BACKUPS, self.backups.len()));
        out.extend_from_slice(&self.backups);
        let warn_blob = self.serialize_warnings();
        out.extend_from_slice(&Self::tag_header(TAG_WARNINGS, warn_blob.len()));
        out.extend_from_slice(&warn_blob);
        // other (omitted)
        Ok(out)
    }

    fn deserialize_ltlv(&mut self, data: &[u8]) -> Result<()> {
        let mut offset = 0usize;
        while offset + 4 <= data.len() {
            let tag = data[offset];
            let len = ((data[offset + 1] as usize) << 16)
                | ((data[offset + 2] as usize) << 8)
                | data[offset + 3] as usize;
            offset += 4;
            if offset + len > data.len() {
                return Err(MegaError::InvalidResponse);
            }
            let slice = &data[offset..offset + len];
            offset += len;
            match tag {
                TAG_VERSION => {
                    if !slice.is_empty() {
                        self.version = slice[0];
                    }
                }
                TAG_CREATION_TIME => {
                    if slice.len() == 4 {
                        self.creation_time = u32::from_be_bytes(slice.try_into().unwrap());
                    }
                }
                TAG_IDENTITY => {
                    if slice.len() == 8 {
                        self.identity = u64::from_le_bytes(slice.try_into().unwrap());
                    }
                }
                TAG_GENERATION => {
                    if slice.len() == 4 {
                        let gen_wire = u32::from_be_bytes(slice.try_into().unwrap());
                        // SDK stores generation+1 on wire; avoid underflow on legacy zero
                        self.generation = gen_wire.saturating_sub(1);
                    }
                }
                TAG_ATTR => self.attr = slice.to_vec(),
                TAG_PRIV_ED25519 => self.priv_ed25519 = slice.to_vec(),
                TAG_PRIV_CU25519 => self.priv_cu25519 = slice.to_vec(),
                TAG_PRIV_RSA => {
                    // SDK expects empty or >=512 bytes (short format); otherwise invalid.
                    if !slice.is_empty() && slice.len() < 512 {
                        return Err(MegaError::InvalidResponse);
                    }
                    self.priv_rsa = slice.to_vec();
                }
                TAG_AUTHRING_ED25519 => {
                    Self::validate_authring(slice)?;
                    self.authring_ed = AuthRing::deserialize_ltlv(slice);
                }
                TAG_AUTHRING_CU25519 => {
                    Self::validate_authring(slice)?;
                    self.authring_cu = AuthRing::deserialize_ltlv(slice);
                }
                TAG_SHAREKEYS => {
                    self.share_keys = Self::parse_share_keys(slice)?;
                    self.rebuild_share_key_index();
                }
                TAG_PENDING_OUTSHARES => self.pending_out = Self::parse_pending_out(slice)?,
                TAG_PENDING_INSHARES => self.pending_in = Self::parse_pending_in(slice)?,
                TAG_BACKUPS => {
                    // Backups are opaque; guard against absurdly large payloads (cap 1 MiB).
                    if slice.len() > 1_048_576 {
                        return Err(MegaError::InvalidResponse);
                    }
                    self.backups = slice.to_vec();
                }
                TAG_WARNINGS => self.warnings = Self::parse_warnings(slice)?,
                _ => {
                    // ignore unknown tags
                }
            }
        }
        Ok(())
    }

    fn serialize_share_keys(&self) -> Vec<u8> {
        let mut out = Vec::new();
        for sk in &self.share_keys {
            out.extend_from_slice(&sk.handle);
            out.extend_from_slice(&sk.key);
            out.push(sk.flags);
        }
        out
    }

    fn parse_share_keys(data: &[u8]) -> Result<Vec<ShareKeyEntry>> {
        let mut out = Vec::new();
        let mut offset = 0usize;
        while offset + 6 + 16 < data.len() {
            let mut h = [0u8; 6];
            h.copy_from_slice(&data[offset..offset + 6]);
            offset += 6;
            let mut k = [0u8; 16];
            k.copy_from_slice(&data[offset..offset + 16]);
            offset += 16;
            let flags = data[offset];
            offset += 1;
            out.push(ShareKeyEntry {
                handle: h,
                key: k,
                flags,
            });
        }
        Ok(out)
    }

    fn serialize_pending_out(&self) -> Vec<u8> {
        let mut out = Vec::new();
        for entry in &self.pending_out {
            match &entry.uid {
                PendingUid::Email(email) => {
                    if email.len() >= 256 {
                        continue;
                    }
                    out.push(email.len() as u8);
                    out.extend_from_slice(&entry.node_handle);
                    out.extend_from_slice(email.as_bytes());
                }
                PendingUid::UserHandle(uh) => {
                    out.push(0);
                    out.extend_from_slice(&entry.node_handle);
                    out.extend_from_slice(uh);
                }
            }
        }
        out
    }

    fn parse_pending_out(data: &[u8]) -> Result<Vec<PendingOutEntry>> {
        let mut out = Vec::new();
        let mut offset = 0usize;
        while offset + 1 + 6 <= data.len() {
            let len = data[offset] as usize;
            offset += 1;
            if offset + 6 > data.len() {
                break;
            }
            let mut h = [0u8; 6];
            h.copy_from_slice(&data[offset..offset + 6]);
            offset += 6;
            if len == 0 {
                if offset + 8 > data.len() {
                    break;
                }
                let mut uh = [0u8; 8];
                uh.copy_from_slice(&data[offset..offset + 8]);
                offset += 8;
                out.push(PendingOutEntry {
                    node_handle: h,
                    uid: PendingUid::UserHandle(uh),
                });
            } else {
                if offset + len > data.len() {
                    break;
                }
                let email = String::from_utf8_lossy(&data[offset..offset + len]).to_string();
                offset += len;
                out.push(PendingOutEntry {
                    node_handle: h,
                    uid: PendingUid::Email(email),
                });
            }
        }
        Ok(out)
    }

    fn serialize_pending_in(&self) -> Vec<u8> {
        // LTLV: [len(tag) tag len(value) value]*
        let mut entries = Vec::new();
        for entry in &self.pending_in {
            let mut value = Vec::new();
            value.extend_from_slice(&entry.user_handle);
            value.extend_from_slice(&entry.share_key);
            entries.push((entry.node_handle_b64.clone(), value));
        }
        Self::serialize_ltlv_map(entries)
    }

    fn parse_pending_in(data: &[u8]) -> Result<Vec<PendingInEntry>> {
        let map = Self::parse_ltlv_map(data)?;
        let mut out = Vec::new();
        for (tag, value) in map {
            // Legacy edge case: SDK used to store pending inshare value as base64 string.
            let decoded_value =
                if !value.is_empty() && value.iter().all(|b| b.is_ascii()) && value.len() > 12 {
                    if let Ok(txt) = std::str::from_utf8(&value) {
                        base64url_decode(txt).unwrap_or(value.clone())
                    } else {
                        value.clone()
                    }
                } else {
                    value.clone()
                };

            if decoded_value.len() < 8 {
                continue;
            }
            let mut uh = [0u8; 8];
            uh.copy_from_slice(&decoded_value[..8]);
            let share_key = decoded_value[8..].to_vec();
            out.push(PendingInEntry {
                node_handle_b64: tag,
                user_handle: uh,
                share_key,
            });
        }
        Ok(out)
    }

    fn serialize_warnings(&self) -> Vec<u8> {
        let items = self
            .warnings
            .0
            .iter()
            .map(|(t, v)| (t.clone(), v.clone()))
            .collect::<Vec<_>>();
        Self::serialize_ltlv_map(items)
    }

    fn parse_warnings(data: &[u8]) -> Result<Warnings> {
        if data.len() > 64 * 1024 {
            return Err(MegaError::InvalidResponse);
        }
        let map = Self::parse_ltlv_map(data)?;
        Ok(Warnings(map))
    }

    fn serialize_ltlv_map(entries: Vec<(String, Vec<u8>)>) -> Vec<u8> {
        let mut out = Vec::new();
        for (tag, value) in entries {
            if tag.len() > 255 {
                continue;
            }
            out.push(tag.len() as u8);
            out.extend_from_slice(tag.as_bytes());
            if value.len() < 0xFFFF {
                let len_be = (value.len() as u16).to_be_bytes();
                out.extend_from_slice(&len_be);
            } else {
                out.extend_from_slice(&0xFFFFu16.to_be_bytes());
                let len_be = (value.len() as u32).to_be_bytes();
                out.extend_from_slice(&len_be);
            }
            out.extend_from_slice(&value);
        }
        out
    }

    fn parse_ltlv_map(data: &[u8]) -> Result<Vec<(String, Vec<u8>)>> {
        let mut out = Vec::new();
        let mut offset = 0usize;
        while offset < data.len() {
            if offset >= data.len() {
                break;
            }
            let tag_len = data[offset] as usize;
            offset += 1;
            if offset + tag_len > data.len() {
                break;
            }
            let tag = String::from_utf8_lossy(&data[offset..offset + tag_len]).to_string();
            offset += tag_len;
            if offset + 2 > data.len() {
                break;
            }
            let len16 = u16::from_be_bytes([data[offset], data[offset + 1]]);
            offset += 2;
            let mut val_len: usize = len16 as usize;
            if len16 == 0xFFFF {
                if offset + 4 > data.len() {
                    break;
                }
                let len32 = u32::from_be_bytes([
                    data[offset],
                    data[offset + 1],
                    data[offset + 2],
                    data[offset + 3],
                ]);
                val_len = len32 as usize;
                offset += 4;
            }
            if offset + val_len > data.len() {
                break;
            }
            let value = data[offset..offset + val_len].to_vec();
            offset += val_len;
            out.push((tag, value));
        }
        Ok(out)
    }

    /// Merge another key manager into this one, preserving existing entries.
    ///
    /// Share keys and pending shares are unioned, while metadata fields are only
    /// copied when the current value is empty or zero. Authrings are merged by
    /// unioning entries, and warnings preserve existing tags.
    ///
    /// # Examples
    /// ```
    /// use megalib::base64::base64url_encode;
    /// use megalib::crypto::KeyManager;
    ///
    /// let mut current = KeyManager::new();
    /// let mut incoming = KeyManager::new();
    ///
    /// let handle = base64url_encode(&[15u8; 6]);
    /// incoming.add_share_key_from_str(&handle, &[9u8; 16]);
    ///
    /// current.merge_from(&incoming);
    /// assert!(current.get_share_key_from_str(&handle).is_some());
    /// ```
    pub fn merge_from(&mut self, other: &KeyManager) {
        // Merge share keys (overwrite key+flags if handle matches, OR together flags).
        for entry in &other.share_keys {
            if let Some(&idx) = self.share_key_index.get(&entry.handle) {
                if let Some(existing) = self.share_keys.get_mut(idx) {
                    existing.key = entry.key;
                    existing.flags |= entry.flags;
                }
            } else {
                let idx = self.share_keys.len();
                self.share_keys.push(entry.clone());
                self.share_key_index.insert(entry.handle, idx);
            }
        }

        // Merge pending out shares without duplicates.
        for entry in &other.pending_out {
            if !self.pending_out.iter().any(|e| {
                e.node_handle == entry.node_handle && matches_pending_uid(&e.uid, &entry.uid)
            }) {
                self.pending_out.push(entry.clone());
            }
        }

        // Merge pending in shares without duplicates.
        for entry in &other.pending_in {
            if !self.pending_in.iter().any(|e| {
                e.node_handle_b64 == entry.node_handle_b64 && e.user_handle == entry.user_handle
            }) {
                self.pending_in.push(entry.clone());
            }
        }

        if self.creation_time == 0 {
            self.creation_time = other.creation_time;
        }
        if self.identity == 0 {
            self.identity = other.identity;
        }
        if self.attr.is_empty() {
            self.attr = other.attr.clone();
        }
        if self.priv_ed25519.is_empty() {
            self.priv_ed25519 = other.priv_ed25519.clone();
        }
        if self.priv_cu25519.is_empty() {
            self.priv_cu25519 = other.priv_cu25519.clone();
        }
        if self.priv_rsa.is_empty() {
            self.priv_rsa = other.priv_rsa.clone();
        }
        // Union authrings to retain all entries
        self.authring_ed.merge_union(&other.authring_ed);
        self.authring_cu.merge_union(&other.authring_cu);
        if self.backups.is_empty() {
            self.backups = other.backups.clone();
        }
        // Merge warnings map: keep existing entries, add missing from other.
        for (tag, val) in &other.warnings.0 {
            if !self.warnings.0.iter().any(|(k, _)| k == tag) {
                self.warnings.0.push((tag.clone(), val.clone()));
            }
        }
        self.manual_verification |= other.manual_verification;

        self.generation = self.generation.max(other.generation);
    }
}

fn matches_pending_uid(a: &PendingUid, b: &PendingUid) -> bool {
    match (a, b) {
        (PendingUid::Email(ae), PendingUid::Email(be)) => ae == be,
        (PendingUid::UserHandle(ah), PendingUid::UserHandle(bh)) => ah == bh,
        _ => false,
    }
}

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

    #[test]
    fn roundtrip_keys_container_uses_hkdf() {
        let master = [3u8; 16];
        let mut km = KeyManager::new();
        km.creation_time = 1234;
        km.identity = 0x1122334455667788;
        km.generation = 1;
        km.priv_ed25519 = vec![9u8; 32];
        km.priv_cu25519 = vec![8u8; 32];
        km.priv_rsa = Vec::new();
        km.authring_ed = AuthRing::deserialize_ltlv(&[1, 2, 3]);
        km.authring_cu = AuthRing::deserialize_ltlv(&[4, 5, 6]);
        km.share_keys.push(ShareKeyEntry {
            handle: [1u8; 6],
            key: [7u8; 16],
            flags: 0b11,
        });

        let blob = km.encode_container(&master).expect("encode");
        assert_eq!(blob[0], 20);
        assert_eq!(blob[1], 0);

        let mut decoded = KeyManager::new();
        decoded
            .decode_container(&blob, &master)
            .expect("decode with derived key");

        assert_eq!(decoded.priv_ed25519, km.priv_ed25519);
        assert_eq!(decoded.priv_cu25519, km.priv_cu25519);
        assert_eq!(decoded.share_keys.len(), 1);
        assert_eq!(decoded.share_keys[0].key, [7u8; 16]);
        assert_eq!(decoded.generation, km.generation);
    }

    #[test]
    fn reject_downgrade_on_decode() {
        let master = [9u8; 16];
        let mut current = KeyManager::new();
        current.generation = 5;
        current.priv_ed25519 = vec![1u8; 32];
        current.priv_cu25519 = vec![2u8; 32];

        let mut older = KeyManager::new();
        older.generation = 2;
        older.priv_ed25519 = vec![1u8; 32];
        older.priv_cu25519 = vec![2u8; 32];
        let blob = older.encode_container(&master).expect("encode older");

        let res = current.decode_container(&blob, &master);
        assert!(res.is_err(), "downgrade should be rejected");
    }

    #[test]
    fn merge_unions_flags_and_pending() {
        let handle_a = base64url_encode(&[0u8; 6]);
        let handle_b = base64url_encode(&[1u8; 6]);
        let handle_c = base64url_encode(&[2u8; 6]);

        let mut km1 = KeyManager::new();
        km1.generation = 1;
        km1.priv_ed25519 = vec![1u8; 32];
        km1.priv_cu25519 = vec![2u8; 32];
        km1.add_share_key_with_flags(&handle_a, &[7u8; 16], true, false);
        km1.add_pending_out_email(&handle_a, "a@example.com");

        let mut km2 = KeyManager::new();
        km2.generation = 3;
        km2.priv_ed25519 = vec![1u8; 32];
        km2.priv_cu25519 = vec![2u8; 32];
        km2.add_share_key_with_flags(&handle_a, &[7u8; 16], false, true);
        km2.add_pending_out_user_handle(&handle_b, &[9u8; 8]);
        km2.add_pending_in(&handle_c, &[5u8; 8], vec![1, 2, 3]);

        km1.merge_from(&km2);

        let flags = km1.share_key_flags(&handle_a).unwrap();
        assert!(flags & SHAREKEY_FLAG_TRUSTED != 0);
        assert!(flags & SHAREKEY_FLAG_IN_USE != 0);
        assert_eq!(km1.pending_out.len(), 2);
        assert_eq!(km1.pending_in.len(), 1);
        assert_eq!(km1.generation, 3);
    }

    #[test]
    fn generation_wire_roundtrip_plus_one() {
        let master = [4u8; 16];
        let mut km = KeyManager::new();
        km.generation = 4;
        km.priv_ed25519 = vec![1u8; 32];
        km.priv_cu25519 = vec![2u8; 32];
        let blob = km.encode_container(&master).expect("encode");

        let mut decoded = KeyManager::new();
        decoded.decode_container(&blob, &master).expect("decode");
        assert_eq!(decoded.generation, 4);
    }
}
