//! Message padding and unpadding for traffic analysis resistance.
//!
//! All messages sent over the wire are padded to one of a fixed set of
//! bucket sizes. This prevents a passive observer (e.g., the relay
//! operator) from inferring message types based on ciphertext length.
//!
//! # Bucket sizes
//!
//! | Bucket | Size (bytes) | Typical content |
//! |--------|-------------|-----------------|
//! | 0 | 256 | Small control messages, acks |
//! | 1 | 512 | Short text messages |
//! | 2 | 1,024 | Text messages, small payloads |
//! | 3 | 2,048 | Medium messages |
//! | 4 | 4,096 | Longer messages, small files |
//! | 5 | 8,192 | MLS Commits for small groups |
//! | 6 | 16,384 | MLS Welcomes for small groups |
//! | 7 | 32,768 | MLS Welcomes for larger groups |
//!
//! Messages larger than the largest bucket (32,768 bytes) must be
//! chunked by the caller before padding. This is enforced at the API
//! level.
//!
//! # Wire format
//!
//! ```text
//! [mask: 4 random bytes][masked_length: u32 LE XOR mask (4 bytes)][payload][random_padding]
//! ```
//!
//! The length prefix is XOR'd with a random mask so the header bytes
//! are indistinguishable from random. Total padded size always equals
//! exactly one of the bucket sizes.
//!
//! # Examples
//!
//! ```
//! use skytale_padding::{pad, unpad};
//!
//! let message = b"hello world";
//! let padded = pad(message).expect("message fits in a bucket");
//! assert_eq!(padded.len(), 256); // smallest bucket that fits
//!
//! let recovered = unpad(&padded).expect("valid padded message");
//! assert_eq!(recovered, message);
//! ```

use rand::RngExt;

/// Errors that can occur during padding or unpadding operations.
///
/// # Examples
///
/// ```
/// use skytale_padding::{pad, PaddingError, MAX_PADDABLE_SIZE};
///
/// let too_large = vec![0xFF; MAX_PADDABLE_SIZE + 1];
/// match pad(&too_large) {
///     Err(PaddingError::InvalidPadding(msg)) => assert!(msg.contains("too large")),
///     other => panic!("expected InvalidPadding, got: {other:?}"),
/// }
/// ```
#[derive(Debug, thiserror::Error)]
pub enum PaddingError {
    #[error("{0}")]
    InvalidPadding(String),
}

/// Fixed bucket sizes in ascending order.
///
/// Every padded message is exactly one of these sizes.
pub const BUCKET_SIZES: [usize; 8] = [256, 512, 1024, 2048, 4096, 8192, 16384, 32768];

/// Size of the header in the padded format (4-byte mask + 4-byte masked length).
const PAD_HEADER_SIZE: usize = 8;

/// Maximum payload that can be padded into fixed buckets (largest bucket minus header).
pub const MAX_PADDABLE_SIZE: usize = BUCKET_SIZES[BUCKET_SIZES.len() - 1] - PAD_HEADER_SIZE;

/// Maximum payload that can be padded with Padme (limited by u32 length field).
pub const MAX_PADME_SIZE: usize = (u32::MAX as usize) - PAD_HEADER_SIZE;

/// Magic byte prefix distinguishing Padme-padded messages from bucket-padded.
///
/// Wire format for Padme:
/// ```text
/// [0xFE][mask: 4 bytes][masked_length: u32 LE XOR mask][payload][random_padding]
/// ```
///
/// Fixed-bucket messages never start with 0xFE because their first byte is
/// a random mask byte — the chance of collision is 1/256 per message, but
/// `unpad_auto` checks total size against bucket sizes first, so there is
/// no ambiguity.
const PADME_MAGIC: u8 = 0xFE;

/// Total header size for Padme-padded messages (1 magic + 4 mask + 4 masked length).
const PADME_HEADER_SIZE: usize = 9;

/// Pads a payload to the next bucket size.
///
/// The padded output has the format:
/// `[mask: 4 random bytes][masked_length: u32 LE XOR mask][payload][random_bytes]`
///
/// The total output length is the smallest bucket size that can
/// contain the 8-byte header plus the payload.
///
/// # Errors
///
/// Returns [`PaddingError::InvalidPadding`] if the payload is too large
/// to fit in any bucket (i.e., `payload.len() > MAX_PADDABLE_SIZE`).
///
/// # Examples
///
/// ```
/// use skytale_padding::{pad, BUCKET_SIZES};
///
/// // 10-byte payload -> 256-byte bucket
/// let padded = pad(b"0123456789").unwrap();
/// assert_eq!(padded.len(), BUCKET_SIZES[0]);
///
/// // Empty payload -> 256-byte bucket
/// let padded = pad(b"").unwrap();
/// assert_eq!(padded.len(), BUCKET_SIZES[0]);
/// ```
pub fn pad(payload: &[u8]) -> Result<Vec<u8>, PaddingError> {
    let needed = PAD_HEADER_SIZE + payload.len();

    let bucket_size = BUCKET_SIZES
        .iter()
        .find(|&&size| size >= needed)
        .ok_or_else(|| {
            PaddingError::InvalidPadding(format!(
                "payload too large for any bucket: {} bytes (max {})",
                payload.len(),
                MAX_PADDABLE_SIZE
            ))
        })?;

    let mut output = Vec::with_capacity(*bucket_size);

    // Generate 4-byte random mask and XOR with length to hide structure.
    let mut mask = [0u8; 4];
    rand::rng().fill(&mut mask);
    let len_bytes = (payload.len() as u32).to_le_bytes();
    let masked_len = [
        len_bytes[0] ^ mask[0],
        len_bytes[1] ^ mask[1],
        len_bytes[2] ^ mask[2],
        len_bytes[3] ^ mask[3],
    ];
    output.extend_from_slice(&mask);
    output.extend_from_slice(&masked_len);

    // Write the actual payload
    output.extend_from_slice(payload);

    // Fill the remainder with random bytes (indistinguishable from
    // ciphertext to an observer)
    let padding_len = bucket_size - output.len();
    let mut random_padding = vec![0u8; padding_len];
    rand::rng().fill(&mut random_padding);
    output.extend_from_slice(&random_padding);

    debug_assert_eq!(output.len(), *bucket_size);

    Ok(output)
}

/// Removes padding and recovers the original payload.
///
/// Reads the 4-byte random mask and 4-byte masked length, XOR's
/// them to recover the original length, then returns the payload.
///
/// # Errors
///
/// Returns [`PaddingError::InvalidPadding`] if:
/// - The input is shorter than 8 bytes (no header).
/// - The declared length exceeds the available data.
/// - The input size does not match any bucket size (corruption).
///
/// # Examples
///
/// ```
/// use skytale_padding::{pad, unpad};
///
/// let original = b"test message";
/// let padded = pad(original).unwrap();
/// let recovered = unpad(&padded).unwrap();
/// assert_eq!(recovered, original);
/// ```
pub fn unpad(padded: &[u8]) -> Result<&[u8], PaddingError> {
    // Validate bucket size
    if !BUCKET_SIZES.contains(&padded.len()) {
        return Err(PaddingError::InvalidPadding(format!(
            "invalid padded message size: {} (expected one of {:?})",
            padded.len(),
            BUCKET_SIZES
        )));
    }

    if padded.len() < PAD_HEADER_SIZE {
        return Err(PaddingError::InvalidPadding(
            "padded message too short for length header".to_string(),
        ));
    }

    // Read 4-byte mask and 4-byte masked length, XOR to recover.
    let mask: [u8; 4] = padded[..4]
        .try_into()
        .map_err(|_| PaddingError::InvalidPadding("failed to read mask".to_string()))?;
    let masked: [u8; 4] = padded[4..8]
        .try_into()
        .map_err(|_| PaddingError::InvalidPadding("failed to read masked length".to_string()))?;
    let len_bytes = [
        masked[0] ^ mask[0],
        masked[1] ^ mask[1],
        masked[2] ^ mask[2],
        masked[3] ^ mask[3],
    ];
    let original_len = u32::from_le_bytes(len_bytes) as usize;

    let available = padded.len() - PAD_HEADER_SIZE;
    if original_len > available {
        return Err(PaddingError::InvalidPadding(format!(
            "declared length {} exceeds available data {}",
            original_len, available
        )));
    }

    Ok(&padded[PAD_HEADER_SIZE..PAD_HEADER_SIZE + original_len])
}

/// Returns the bucket size that a payload of the given length would
/// be padded to.
///
/// Returns `None` if the payload is too large for any bucket.
///
/// # Examples
///
/// ```
/// use skytale_padding::bucket_for_size;
///
/// assert_eq!(bucket_for_size(10), Some(256));
/// assert_eq!(bucket_for_size(248), Some(256));   // 248 + 8 byte header = 256
/// assert_eq!(bucket_for_size(249), Some(512));   // 249 + 8 = 257 -> next bucket
/// assert_eq!(bucket_for_size(50000), None);
/// ```
pub fn bucket_for_size(payload_len: usize) -> Option<usize> {
    let needed = PAD_HEADER_SIZE + payload_len;
    BUCKET_SIZES.iter().find(|&&size| size >= needed).copied()
}

/// Computes the Padme padded length for a given message length.
///
/// The Padme algorithm (from the PURB paper, PETS 2019) rounds up
/// message lengths such that information leakage is O(log(log(M)))
/// bits, with at most ~12% overhead.
///
/// # Algorithm
///
/// For length L:
/// 1. E = floor(log2(L))
/// 2. S = floor(log2(E)) + 1
/// 3. last_bits = E - S
/// 4. bit_mask = (1 << last_bits) - 1
/// 5. padded = (L + bit_mask) & !bit_mask
///
/// # Examples
///
/// ```
/// use skytale_padding::padme_length;
///
/// assert_eq!(padme_length(1), 2);
/// assert_eq!(padme_length(9), 10);
/// assert_eq!(padme_length(100), 104);
/// assert_eq!(padme_length(1000), 1024);
/// ```
pub fn padme_length(len: usize) -> usize {
    if len <= 2 {
        return 2;
    }

    // floor(log2(n)) for positive n
    let ilog2 = |n: u32| -> u32 { u32::BITS - 1 - n.leading_zeros() };

    // E = floor(log2(len))
    let e = ilog2(len as u32);

    // S = floor(log2(E)) + 1
    let s = if e <= 1 { 1 } else { ilog2(e) + 1 };

    // Number of trailing bits to zero out
    let last_bits = e.saturating_sub(s);
    let bit_mask = (1usize << last_bits) - 1;

    // Round up by adding bit_mask then masking
    (len + bit_mask) & !bit_mask
}

/// Pads a payload using the Padme algorithm for large messages.
///
/// Uses the same mask+XOR wire format as fixed-bucket padding but
/// with a `0xFE` magic byte prefix and Padme-determined output size.
/// This handles payloads of any size up to [`MAX_PADME_SIZE`].
///
/// # Wire format
///
/// ```text
/// [0xFE][mask: 4 bytes][masked_length: u32 LE XOR mask][payload][random_padding]
/// ```
///
/// # Errors
///
/// Returns [`PaddingError::InvalidPadding`] if the payload exceeds
/// [`MAX_PADME_SIZE`].
///
/// # Examples
///
/// ```
/// use skytale_padding::{pad_large, unpad_large};
///
/// let message = vec![0xAB; 50_000];
/// let padded = pad_large(&message).unwrap();
/// assert!(padded.len() > 50_000);
/// assert!(padded.len() <= 57_000); // Padme: <=12% overhead
/// let recovered = unpad_large(&padded).unwrap();
/// assert_eq!(recovered, &message[..]);
/// ```
pub fn pad_large(payload: &[u8]) -> Result<Vec<u8>, PaddingError> {
    if payload.len() > MAX_PADME_SIZE {
        return Err(PaddingError::InvalidPadding(format!(
            "payload too large for Padme padding: {} bytes (max {})",
            payload.len(),
            MAX_PADME_SIZE
        )));
    }

    let total_needed = PADME_HEADER_SIZE + payload.len();
    let padded_size = padme_length(total_needed);

    let mut output = Vec::with_capacity(padded_size);

    // Magic byte
    output.push(PADME_MAGIC);

    // Generate 4-byte random mask and XOR with length
    let mut mask = [0u8; 4];
    rand::rng().fill(&mut mask);
    let len_bytes = (payload.len() as u32).to_le_bytes();
    let masked_len = [
        len_bytes[0] ^ mask[0],
        len_bytes[1] ^ mask[1],
        len_bytes[2] ^ mask[2],
        len_bytes[3] ^ mask[3],
    ];
    output.extend_from_slice(&mask);
    output.extend_from_slice(&masked_len);

    // Payload
    output.extend_from_slice(payload);

    // Random padding to reach Padme target size
    let padding_len = padded_size - output.len();
    if padding_len > 0 {
        let mut random_padding = vec![0u8; padding_len];
        rand::rng().fill(&mut random_padding);
        output.extend_from_slice(&random_padding);
    }

    debug_assert_eq!(output.len(), padded_size);
    Ok(output)
}

/// Removes Padme padding and recovers the original payload.
///
/// Expects the `0xFE` magic byte prefix followed by the mask+XOR
/// length header.
///
/// # Errors
///
/// Returns [`PaddingError::InvalidPadding`] if:
/// - The input is too short for the header.
/// - The magic byte is not `0xFE`.
/// - The declared length exceeds available data.
///
/// # Examples
///
/// ```
/// use skytale_padding::{pad_large, unpad_large};
///
/// let original = vec![0xCD; 100_000];
/// let padded = pad_large(&original).unwrap();
/// let recovered = unpad_large(&padded).unwrap();
/// assert_eq!(recovered, &original[..]);
/// ```
pub fn unpad_large(padded: &[u8]) -> Result<&[u8], PaddingError> {
    if padded.len() < PADME_HEADER_SIZE {
        return Err(PaddingError::InvalidPadding(
            "Padme-padded message too short for header".to_string(),
        ));
    }

    if padded[0] != PADME_MAGIC {
        return Err(PaddingError::InvalidPadding(format!(
            "expected Padme magic byte 0x{:02X}, got 0x{:02X}",
            PADME_MAGIC, padded[0]
        )));
    }

    let mask: [u8; 4] = padded[1..5]
        .try_into()
        .map_err(|_| PaddingError::InvalidPadding("failed to read mask".to_string()))?;
    let masked: [u8; 4] = padded[5..9]
        .try_into()
        .map_err(|_| PaddingError::InvalidPadding("failed to read masked length".to_string()))?;
    let len_bytes = [
        masked[0] ^ mask[0],
        masked[1] ^ mask[1],
        masked[2] ^ mask[2],
        masked[3] ^ mask[3],
    ];
    let original_len = u32::from_le_bytes(len_bytes) as usize;

    let available = padded.len() - PADME_HEADER_SIZE;
    if original_len > available {
        return Err(PaddingError::InvalidPadding(format!(
            "Padme: declared length {} exceeds available data {}",
            original_len, available
        )));
    }

    Ok(&padded[PADME_HEADER_SIZE..PADME_HEADER_SIZE + original_len])
}

/// Pads a payload using the best strategy for its size.
///
/// - Payloads that fit in fixed buckets (<=32,760 bytes) use bucket padding.
/// - Larger payloads use Padme padding with a `0xFE` prefix.
///
/// Use [`unpad_auto`] to recover the original payload regardless of
/// which strategy was used.
///
/// # Errors
///
/// Returns [`PaddingError::InvalidPadding`] if the payload exceeds
/// [`MAX_PADME_SIZE`].
///
/// # Examples
///
/// ```
/// use skytale_padding::{pad_auto, unpad_auto, BUCKET_SIZES};
///
/// // Small message -> fixed bucket
/// let small = pad_auto(b"hello").unwrap();
/// assert_eq!(small.len(), BUCKET_SIZES[0]);
///
/// // Large message -> Padme
/// let large_msg = vec![0u8; 50_000];
/// let large = pad_auto(&large_msg).unwrap();
/// assert!(large.len() > 50_000);
///
/// // Both round-trip through unpad_auto
/// assert_eq!(unpad_auto(&small).unwrap(), b"hello");
/// assert_eq!(unpad_auto(&large).unwrap(), &large_msg[..]);
/// ```
pub fn pad_auto(payload: &[u8]) -> Result<Vec<u8>, PaddingError> {
    if bucket_for_size(payload.len()).is_some() {
        pad(payload)
    } else {
        pad_large(payload)
    }
}

/// Removes padding (either fixed-bucket or Padme) and recovers the payload.
///
/// Detects the padding strategy by checking if the input size matches a
/// fixed bucket. If it does, uses bucket unpadding. Otherwise, checks for
/// the Padme magic byte and uses Padme unpadding.
///
/// # Errors
///
/// Returns [`PaddingError::InvalidPadding`] if the format is invalid.
///
/// # Examples
///
/// ```
/// use skytale_padding::{pad_auto, unpad_auto};
///
/// let original = b"test message";
/// let padded = pad_auto(original).unwrap();
/// let recovered = unpad_auto(&padded).unwrap();
/// assert_eq!(recovered, original);
/// ```
pub fn unpad_auto(padded: &[u8]) -> Result<&[u8], PaddingError> {
    if BUCKET_SIZES.contains(&padded.len()) {
        unpad(padded)
    } else if padded.first() == Some(&PADME_MAGIC) {
        unpad_large(padded)
    } else {
        Err(PaddingError::InvalidPadding(format!(
            "unknown padding format: size {} is not a bucket and no Padme magic byte",
            padded.len()
        )))
    }
}

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

    #[test]
    fn test_pad_unpad_roundtrip_empty() {
        let original = b"";
        let padded = pad(original).expect("empty payload should pad");
        assert_eq!(padded.len(), 256);
        let recovered = unpad(&padded).expect("should unpad");
        assert_eq!(recovered, original);
    }

    #[test]
    fn test_pad_unpad_roundtrip_small() {
        let original = b"hello, world!";
        let padded = pad(original).expect("should pad");
        assert_eq!(padded.len(), 256);
        let recovered = unpad(&padded).expect("should unpad");
        assert_eq!(recovered, original);
    }

    #[test]
    fn test_pad_unpad_roundtrip_medium() {
        // 300 bytes of payload -> needs 308 with header -> bucket 512
        let original = vec![0xAB; 300];
        let padded = pad(&original).expect("should pad");
        assert_eq!(padded.len(), 512);
        let recovered = unpad(&padded).expect("should unpad");
        assert_eq!(recovered, original);
    }

    #[test]
    fn test_pad_unpad_roundtrip_large() {
        // 2000 bytes -> needs 2008 -> bucket 2048
        let original = vec![0xCD; 2000];
        let padded = pad(&original).expect("should pad");
        assert_eq!(padded.len(), 2048);
        let recovered = unpad(&padded).expect("should unpad");
        assert_eq!(recovered, original);
    }

    #[test]
    fn test_pad_unpad_roundtrip_max_bucket() {
        // Fill the largest bucket exactly (32768 - 8 = 32760 bytes)
        let original = vec![0xEF; MAX_PADDABLE_SIZE];
        let padded = pad(&original).expect("should pad");
        assert_eq!(padded.len(), 32768);
        let recovered = unpad(&padded).expect("should unpad");
        assert_eq!(recovered, original);
    }

    #[test]
    fn test_pad_too_large() {
        let original = vec![0xFF; MAX_PADDABLE_SIZE + 1];
        let result = pad(&original);
        assert!(result.is_err());
        match result {
            Err(PaddingError::InvalidPadding(msg)) => {
                assert!(msg.contains("too large"));
            }
            other => panic!("expected PaddingError, got: {other:?}"),
        }
    }

    #[test]
    fn test_unpad_wrong_size() {
        // 100 bytes is not a valid bucket size
        let garbage = vec![0u8; 100];
        let result = unpad(&garbage);
        assert!(result.is_err());
    }

    #[test]
    fn test_unpad_corrupted_length() {
        // Bucket-sized buffer but declared length exceeds capacity.
        // Format: [mask: 4][masked_len: 4][...]. Use mask=0 so
        // masked_len == raw len for easy construction.
        let mut padded = vec![0u8; 256];
        // mask = [0,0,0,0], masked_len = 300 LE -> decoded len = 300
        // (exceeds 256 - 8 = 248)
        padded[4..8].copy_from_slice(&300u32.to_le_bytes());
        let result = unpad(&padded);
        assert!(result.is_err());
    }

    #[test]
    fn test_bucket_boundaries() {
        // Exactly fits bucket 0: 248 bytes payload + 8 byte header = 256
        let boundary = vec![0u8; 248];
        let padded = pad(&boundary).expect("should fit in bucket 0");
        assert_eq!(padded.len(), 256);

        // One byte over: 249 bytes payload + 8 byte header = 257 -> bucket 512
        let over_boundary = vec![0u8; 249];
        let padded = pad(&over_boundary).expect("should fit in bucket 1");
        assert_eq!(padded.len(), 512);
    }

    #[test]
    fn test_bucket_for_size_values() {
        // Header is 8 bytes, so boundary = bucket_size - 8
        assert_eq!(bucket_for_size(0), Some(256));
        assert_eq!(bucket_for_size(1), Some(256));
        assert_eq!(bucket_for_size(248), Some(256)); // 248 + 8 = 256
        assert_eq!(bucket_for_size(249), Some(512)); // 249 + 8 = 257
        assert_eq!(bucket_for_size(504), Some(512)); // 504 + 8 = 512
        assert_eq!(bucket_for_size(505), Some(1024)); // 505 + 8 = 513
        assert_eq!(bucket_for_size(1016), Some(1024)); // 1016 + 8 = 1024
        assert_eq!(bucket_for_size(1017), Some(2048)); // 1017 + 8 = 1025
        assert_eq!(bucket_for_size(2040), Some(2048)); // 2040 + 8 = 2048
        assert_eq!(bucket_for_size(2041), Some(4096)); // 2041 + 8 = 2049
        assert_eq!(bucket_for_size(4088), Some(4096)); // 4088 + 8 = 4096
        assert_eq!(bucket_for_size(4089), Some(8192)); // 4089 + 8 = 4097
        assert_eq!(bucket_for_size(8184), Some(8192)); // 8184 + 8 = 8192
        assert_eq!(bucket_for_size(8185), Some(16384)); // 8185 + 8 = 8193
        assert_eq!(bucket_for_size(16376), Some(16384)); // 16376 + 8 = 16384
        assert_eq!(bucket_for_size(16377), Some(32768)); // 16377 + 8 = 16385
        assert_eq!(bucket_for_size(32760), Some(32768)); // 32760 + 8 = 32768
        assert_eq!(bucket_for_size(32761), None); // 32761 + 8 = 32769
    }

    // --- Padme algorithm tests ---

    #[test]
    fn test_padme_length_small_values() {
        assert_eq!(padme_length(1), 2);
        assert_eq!(padme_length(2), 2);
        // E=1, S=1, last_bits=0 -> identity for 3-7
        assert_eq!(padme_length(3), 3);
        assert_eq!(padme_length(4), 4);
        assert_eq!(padme_length(5), 5);
        assert_eq!(padme_length(7), 7);
        // E=3, S=2, last_bits=1: rounding starts at 8+
        assert_eq!(padme_length(8), 8);
        assert_eq!(padme_length(9), 10);
        assert_eq!(padme_length(15), 16);
    }

    #[test]
    fn test_padme_length_medium_values() {
        // 100 bytes: E=6, S=3, last_bits=3, mask=7 -> (100+7)&!7 = 104
        assert_eq!(padme_length(100), 104);
        // 1000 bytes: E=9, S=4, last_bits=5, mask=31 -> (1000+31)&!31 = 1024
        assert_eq!(padme_length(1000), 1024);
    }

    #[test]
    fn test_padme_length_large_values() {
        let padded = padme_length(50_000);
        assert!(padded >= 50_000);
        // Padme guarantees at most ~12% overhead
        assert!(padded <= 56_000, "overhead too high: {padded}");
    }

    #[test]
    fn test_padme_length_always_geq_input() {
        for size in [1, 2, 3, 10, 100, 1000, 10_000, 100_000, 1_000_000] {
            let padded = padme_length(size);
            assert!(padded >= size, "padme_length({size}) = {padded} < {size}");
        }
    }

    #[test]
    fn test_padme_overhead_bounded() {
        // For messages >= 256 bytes, overhead should be <= 12%
        for &size in &[256, 512, 1000, 5000, 50_000, 500_000, 5_000_000] {
            let padded = padme_length(size);
            let overhead_pct = ((padded - size) as f64 / size as f64) * 100.0;
            assert!(
                overhead_pct <= 13.0,
                "padme_length({size}) = {padded}, overhead {overhead_pct:.1}% > 12%"
            );
        }
    }

    // --- Padme pad/unpad tests ---

    #[test]
    fn test_pad_large_unpad_large_roundtrip() {
        let original = vec![0xAB; 50_000];
        let padded = pad_large(&original).expect("should pad");
        assert!(padded[0] == PADME_MAGIC);
        let recovered = unpad_large(&padded).expect("should unpad");
        assert_eq!(recovered, &original[..]);
    }

    #[test]
    fn test_pad_large_unpad_large_100kb() {
        let original = vec![0xCD; 100_000];
        let padded = pad_large(&original).expect("should pad");
        let recovered = unpad_large(&padded).expect("should unpad");
        assert_eq!(recovered, &original[..]);
    }

    #[test]
    fn test_pad_large_unpad_large_1mb() {
        let original = vec![0xEF; 1_000_000];
        let padded = pad_large(&original).expect("should pad");
        let recovered = unpad_large(&padded).expect("should unpad");
        assert_eq!(recovered, &original[..]);
    }

    #[test]
    fn test_pad_large_small_payload() {
        // Padme can handle small payloads too (though fixed buckets are preferred)
        let original = b"hello";
        let padded = pad_large(original).expect("should pad");
        let recovered = unpad_large(&padded).expect("should unpad");
        assert_eq!(recovered, original);
    }

    #[test]
    fn test_unpad_large_wrong_magic() {
        let mut padded = pad_large(b"test").expect("should pad");
        padded[0] = 0xFF; // wrong magic
        assert!(unpad_large(&padded).is_err());
    }

    #[test]
    fn test_unpad_large_too_short() {
        let short = vec![PADME_MAGIC; 5]; // too short for header
        assert!(unpad_large(&short).is_err());
    }

    // --- Auto pad/unpad tests ---

    #[test]
    fn test_pad_auto_small_uses_buckets() {
        let padded = pad_auto(b"hello").expect("should pad");
        assert_eq!(padded.len(), 256); // bucket 0
        assert_ne!(padded[0], PADME_MAGIC); // not Padme (well, might be random)
                                            // But unpad_auto should work because size matches a bucket
        let recovered = unpad_auto(&padded).expect("should unpad");
        assert_eq!(recovered, b"hello");
    }

    #[test]
    fn test_pad_auto_large_uses_padme() {
        let original = vec![0xAB; 50_000];
        let padded = pad_auto(&original).expect("should pad");
        assert!(!BUCKET_SIZES.contains(&padded.len())); // not a bucket size
        assert_eq!(padded[0], PADME_MAGIC); // Padme magic
        let recovered = unpad_auto(&padded).expect("should unpad");
        assert_eq!(recovered, &original[..]);
    }

    #[test]
    fn test_pad_auto_boundary_uses_bucket() {
        // Exactly MAX_PADDABLE_SIZE should use bucket padding
        let original = vec![0u8; MAX_PADDABLE_SIZE];
        let padded = pad_auto(&original).expect("should pad");
        assert_eq!(padded.len(), 32768); // largest bucket
        let recovered = unpad_auto(&padded).expect("should unpad");
        assert_eq!(recovered, &original[..]);
    }

    #[test]
    fn test_pad_auto_over_boundary_uses_padme() {
        // One byte over MAX_PADDABLE_SIZE should use Padme
        let original = vec![0u8; MAX_PADDABLE_SIZE + 1];
        let padded = pad_auto(&original).expect("should pad");
        assert!(!BUCKET_SIZES.contains(&padded.len()));
        assert_eq!(padded[0], PADME_MAGIC);
        let recovered = unpad_auto(&padded).expect("should unpad");
        assert_eq!(recovered, &original[..]);
    }

    #[test]
    fn test_unpad_auto_unknown_format() {
        // Not a bucket size, no Padme magic
        let garbage = vec![0x42; 100];
        assert!(unpad_auto(&garbage).is_err());
    }

    #[test]
    fn test_padding_is_random() {
        // Pad the same message twice; the entire output should differ
        // because the mask and padding are random.
        let original = b"deterministic payload";
        let padded_a = pad(original).expect("should pad");
        let padded_b = pad(original).expect("should pad");

        // Both must roundtrip correctly
        assert_eq!(unpad(&padded_a).unwrap(), original);
        assert_eq!(unpad(&padded_b).unwrap(), original);

        // The header bytes (mask + masked_len) should differ because
        // each call generates a fresh random mask.
        assert_ne!(
            padded_a[..PAD_HEADER_SIZE],
            padded_b[..PAD_HEADER_SIZE],
            "header should differ due to random mask"
        );

        // The overall padded output should differ
        assert_ne!(padded_a, padded_b, "padded outputs should differ");
    }
}
