//! PDF encryption algorithms.
//!
//! This module implements the cryptographic algorithms specified in the PDF specification
//! for key derivation and password validation.
//!
//! PDF Spec: Section 7.6.3 - Standard Security Handler
//! PDF 2.0 Spec (ISO 32000-2:2020): Section 7.6.4.3.3 - Algorithm 8-11 for R>=5

use md5::{Digest, Md5};
use sha2::Sha256;

/// Padding string used in PDF encryption (32 bytes).
///
/// PDF Spec: Algorithm 2, step 1
const PADDING: &[u8; 32] = b"\x28\xBF\x4E\x5E\x4E\x75\x8A\x41\
                              \x64\x00\x4E\x56\xFF\xFA\x01\x08\
                              \x2E\x2E\x00\xB6\xD0\x68\x3E\x80\
                              \x2F\x0C\xA9\xFE\x64\x53\x69\x7A";

/// Compute the encryption key from a password (Algorithm 2).
///
/// PDF Spec: Section 7.6.3.3 - Algorithm 2: Computing an encryption key
///
/// # Arguments
///
/// * `password` - User or owner password (up to 32 bytes)
/// * `owner_key` - 32-byte owner password hash from encryption dictionary
/// * `permissions` - User access permissions (P field)
/// * `file_id` - First element of file identifier array
/// * `revision` - Encryption revision number (R field)
/// * `key_length` - Key length in bytes
/// * `encrypt_metadata` - Whether to encrypt metadata
///
/// # Returns
///
/// The derived encryption key
pub fn compute_encryption_key(
    password: &[u8],
    owner_key: &[u8],
    permissions: i32,
    file_id: &[u8],
    revision: u32,
    key_length: usize,
    encrypt_metadata: bool,
) -> Vec<u8> {
    // For R>=5, the encryption key is randomly generated, not derived from password
    // PDF 2.0 Spec: Algorithm 8 generates a random 32-byte file encryption key
    if revision >= 5 {
        return generate_random_encryption_key(key_length);
    }

    let mut hasher = Md5::new();

    // Step a: Pad or truncate password to 32 bytes
    let mut padded_password = [0u8; 32];
    let pass_len = password.len().min(32);
    padded_password[..pass_len].copy_from_slice(&password[..pass_len]);
    if pass_len < 32 {
        padded_password[pass_len..].copy_from_slice(&PADDING[..(32 - pass_len)]);
    }

    // Step b: Pass the password to MD5
    hasher.update(padded_password);

    // Step c: Pass the owner password hash
    hasher.update(owner_key);

    // Step d: Pass permissions as 32-bit little-endian
    hasher.update(permissions.to_le_bytes());

    // Step e: Pass the file identifier
    hasher.update(file_id);

    // Step f: For R >= 4, if EncryptMetadata is false, pass 0xFFFFFFFF
    if revision >= 4 && !encrypt_metadata {
        hasher.update([0xFF, 0xFF, 0xFF, 0xFF]);
    }

    // Step g: Finish MD5 hash
    let mut hash = hasher.finalize().to_vec();

    // Step h: For R >= 3, do 50 additional MD5 iterations on first key_length bytes
    if revision >= 3 {
        for _ in 0..50 {
            let mut hasher = Md5::new();
            hasher.update(&hash[..key_length.min(16)]);
            hash = hasher.finalize().to_vec();
        }
    }

    // Step i: Return first key_length bytes (max 16 for MD5)
    hash[..key_length.min(16)].to_vec()
}

/// Generate a random encryption key for R>=5.
///
/// PDF 2.0 Spec: For AES-256, the file encryption key is randomly generated.
fn generate_random_encryption_key(key_length: usize) -> Vec<u8> {
    use sha2::{Digest, Sha256};

    // Generate random bytes using multiple UUID/timestamp combinations
    let mut key = Vec::with_capacity(key_length);

    while key.len() < key_length {
        let uuid1 = uuid::Uuid::new_v4();
        let uuid2 = uuid::Uuid::new_v4();

        let mut hasher = Sha256::new();
        hasher.update(uuid1.as_bytes());
        hasher.update(uuid2.as_bytes());

        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default();
        hasher.update(now.as_nanos().to_le_bytes());

        let hash = hasher.finalize();
        let remaining = key_length - key.len();
        key.extend_from_slice(&hash[..remaining.min(32)]);
    }

    key
}

/// Pad or truncate a password to 32 bytes using the standard padding.
///
/// PDF Spec: Algorithm 2, step 1
#[allow(dead_code)]
pub fn pad_password(password: &[u8]) -> Vec<u8> {
    let mut padded = Vec::with_capacity(32);
    let pass_len = password.len().min(32);
    padded.extend_from_slice(&password[..pass_len]);
    if pass_len < 32 {
        padded.extend_from_slice(&PADDING[..(32 - pass_len)]);
    }
    padded
}

/// Authenticate the user password (Algorithm 4/5).
///
/// PDF Spec: Section 7.6.3.4 - Algorithm 4/5: User password authentication
///
/// Returns the encryption key if authentication succeeds.
pub fn authenticate_user_password(
    password: &[u8],
    user_key: &[u8],
    owner_key: &[u8],
    permissions: i32,
    file_id: &[u8],
    revision: u32,
    key_length: usize,
    encrypt_metadata: bool,
) -> Option<Vec<u8>> {
    // Compute encryption key from password
    let key = compute_encryption_key(
        password,
        owner_key,
        permissions,
        file_id,
        revision,
        key_length,
        encrypt_metadata,
    );

    // Compute expected user key
    let expected_user_key = if revision >= 3 {
        compute_user_key_r3(&key, file_id)
    } else {
        compute_user_key_r2(&key)
    };

    // Compare first 16 bytes (constant-time comparison)
    let matches = constant_time_compare(&user_key[..16], &expected_user_key[..16]);

    if matches {
        Some(key)
    } else {
        None
    }
}

/// Compute the user password hash for R=2 (Algorithm 4).
///
/// PDF Spec: Section 7.6.3.4 - Algorithm 4
fn compute_user_key_r2(key: &[u8]) -> Vec<u8> {
    // Encrypt padding string with key
    super::rc4::rc4_crypt(key, PADDING)
}

/// Compute the user password hash for R>=3 (Algorithm 5).
///
/// PDF Spec: Section 7.6.3.4 - Algorithm 5
fn compute_user_key_r3(key: &[u8], file_id: &[u8]) -> Vec<u8> {
    // Step a: Create MD5 hash of padding + file ID
    let mut hasher = Md5::new();
    hasher.update(PADDING);
    hasher.update(file_id);
    let mut hash = hasher.finalize().to_vec();

    // Step b: Encrypt the hash 20 times with modified keys
    for i in 0..20 {
        let mut modified_key = key.to_vec();
        for byte in &mut modified_key {
            *byte ^= i as u8;
        }
        hash = super::rc4::rc4_crypt(&modified_key, &hash);
    }

    // Step c: Append 16 arbitrary bytes (we use zeros)
    hash.extend_from_slice(&[0u8; 16]);
    hash
}

/// Compute the owner password hash (Algorithm 3 for R<=4, Algorithm 8 for R>=5).
///
/// PDF Spec: Section 7.6.3.3 - Algorithm 3: Computing the O value (R<=4)
/// PDF 2.0 Spec: Algorithm 8: Computing O and U for R>=5
///
/// This generates the /O value for the encryption dictionary.
///
/// # Arguments
///
/// * `owner_password` - Owner password (if empty, uses user_password)
/// * `user_password` - User password
/// * `revision` - Encryption revision (R value: 2, 3, 4, 5, or 6)
/// * `key_length` - Key length in bytes (5 for 40-bit, 16 for 128-bit, 32 for 256-bit)
///
/// # Returns
///
/// 32-byte owner password hash for /O entry (48 bytes for R>=5)
pub fn compute_owner_password_hash(
    owner_password: &[u8],
    user_password: &[u8],
    revision: u32,
    key_length: usize,
) -> Vec<u8> {
    // For R>=5, use SHA-256 based algorithm (Algorithm 8)
    if revision >= 5 {
        return compute_owner_hash_r5(owner_password, user_password);
    }

    // Algorithm 3 for R<=4
    // Step a: Use owner password, or user password if owner is empty
    let password = if owner_password.is_empty() {
        user_password
    } else {
        owner_password
    };

    // Step b: Pad the password to 32 bytes
    let padded_password = pad_password(password);

    // Step c: Initialize MD5 and pass the padded password
    let mut hasher = Md5::new();
    hasher.update(&padded_password);
    let mut hash = hasher.finalize().to_vec();

    // Step d: For R >= 3, do 50 additional MD5 iterations
    if revision >= 3 {
        for _ in 0..50 {
            let mut hasher = Md5::new();
            hasher.update(&hash[..key_length.min(16)]);
            hash = hasher.finalize().to_vec();
        }
    }

    // Step e: Use first key_length bytes as RC4 key (max 16)
    let rc4_key_len = key_length.min(16);
    let rc4_key = &hash[..rc4_key_len];

    // Step f: Pad the user password
    let padded_user = pad_password(user_password);

    // Step g: RC4 encrypt the padded user password
    let mut result = super::rc4::rc4_crypt(rc4_key, &padded_user);

    // Step h: For R >= 3, do 19 more RC4 encryptions with XOR'd keys
    if revision >= 3 {
        for i in 1..=19 {
            let mut modified_key = rc4_key.to_vec();
            for byte in &mut modified_key {
                *byte ^= i as u8;
            }
            result = super::rc4::rc4_crypt(&modified_key, &result);
        }
    }

    result
}

/// Compute owner password hash for R>=5 (Algorithm 8 part).
///
/// PDF 2.0 Spec: Algorithm 8 - Computing O value for R=5/6
///
/// For R>=5, the O value is 48 bytes:
/// - Bytes 0-31: SHA-256(password || owner_validation_salt || U[0..48])
/// - Bytes 32-39: owner_validation_salt (random 8 bytes)
/// - Bytes 40-47: owner_key_salt (random 8 bytes)
fn compute_owner_hash_r5(owner_password: &[u8], _user_password: &[u8]) -> Vec<u8> {
    // Generate random salts
    let validation_salt = generate_random_bytes(8);
    let key_salt = generate_random_bytes(8);

    // For the initial O computation, we don't have U yet, so we compute a placeholder
    // In practice, the EncryptDictBuilder computes U first, then O
    // For now, compute without U (this is a simplified version)
    let password = truncate_password_utf8(owner_password);

    let mut hasher = Sha256::new();
    hasher.update(&password);
    hasher.update(&validation_salt);
    // Note: In full implementation, we'd include U[0..48] here
    let hash = hasher.finalize();

    // Build 48-byte O value
    let mut result = hash.to_vec(); // 32 bytes
    result.extend_from_slice(&validation_salt); // 8 bytes
    result.extend_from_slice(&key_salt); // 8 bytes

    result
}

/// Compute the user password hash for the encryption dictionary (Algorithm 4/5/8).
///
/// PDF Spec: Section 7.6.3.4 - Algorithm 4 (R=2) and Algorithm 5 (R>=3)
/// PDF 2.0 Spec: Algorithm 8 - Computing U for R>=5
///
/// This generates the /U value for the encryption dictionary.
///
/// # Arguments
///
/// * `encryption_key` - The computed encryption key from Algorithm 2
/// * `file_id` - First element of file identifier array (only used for R>=3)
/// * `revision` - Encryption revision (R value)
///
/// # Returns
///
/// 32-byte user password hash for /U entry (48 bytes for R>=5)
pub fn compute_user_password_hash(encryption_key: &[u8], file_id: &[u8], revision: u32) -> Vec<u8> {
    if revision >= 5 {
        // For R>=5, use the encryption key directly as user password indicator
        // This creates the U value with validation/key salts
        compute_user_hash_r5(encryption_key)
    } else if revision >= 3 {
        compute_user_key_r3(encryption_key, file_id)
    } else {
        compute_user_key_r2(encryption_key)
    }
}

/// Compute user password hash for R>=5 (Algorithm 8 part).
///
/// PDF 2.0 Spec: Algorithm 8 - Computing U value for R=5/6
///
/// For R>=5, the U value is 48 bytes:
/// - Bytes 0-31: SHA-256(password || user_validation_salt)
/// - Bytes 32-39: user_validation_salt (random 8 bytes)
/// - Bytes 40-47: user_key_salt (random 8 bytes)
fn compute_user_hash_r5(user_password: &[u8]) -> Vec<u8> {
    let validation_salt = generate_random_bytes(8);
    let key_salt = generate_random_bytes(8);

    let password = truncate_password_utf8(user_password);

    let mut hasher = Sha256::new();
    hasher.update(&password);
    hasher.update(&validation_salt);
    let hash = hasher.finalize();

    // Build 48-byte U value
    let mut result = hash.to_vec(); // 32 bytes
    result.extend_from_slice(&validation_salt); // 8 bytes
    result.extend_from_slice(&key_salt); // 8 bytes

    result
}

/// Generate random bytes using UUID v4 and timestamp mixing.
fn generate_random_bytes(len: usize) -> Vec<u8> {
    use md5::{Digest, Md5};

    let mut result = Vec::with_capacity(len);

    while result.len() < len {
        let uuid = uuid::Uuid::new_v4();
        let mut hasher = Md5::new();
        hasher.update(uuid.as_bytes());

        // Add timestamp for extra entropy
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default();
        hasher.update(now.as_nanos().to_le_bytes());

        let hash = hasher.finalize();
        let remaining = len - result.len();
        result.extend_from_slice(&hash[..remaining.min(16)]);
    }

    result
}

/// Truncate password to 127 bytes for UTF-8 (R>=5 requirement).
///
/// PDF 2.0 Spec: For R>=5, passwords are UTF-8 encoded and
/// limited to 127 bytes.
fn truncate_password_utf8(password: &[u8]) -> Vec<u8> {
    let mut result = password.to_vec();
    if result.len() > 127 {
        // Find UTF-8 boundary for truncation
        let mut end = 127;
        while end > 0 && (result[end] & 0xC0) == 0x80 {
            end -= 1;
        }
        result.truncate(end);
    }
    result
}

/// Constant-time comparison to prevent timing attacks.
///
/// Returns true if the slices are equal.
fn constant_time_compare(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }

    let mut result = 0u8;
    for (x, y) in a.iter().zip(b.iter()) {
        result |= x ^ y;
    }

    result == 0
}

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

    #[test]
    fn test_pad_password() {
        let password = b"test";
        let padded = pad_password(password);
        assert_eq!(padded.len(), 32);
        assert_eq!(&padded[..4], b"test");
        assert_eq!(&padded[4..], &PADDING[..28]);
    }

    #[test]
    fn test_pad_password_long() {
        let password = b"this is a very long password that exceeds 32 bytes";
        let padded = pad_password(password);
        assert_eq!(padded.len(), 32);
        assert_eq!(&padded[..], &password[..32]);
    }

    #[test]
    fn test_pad_password_exact() {
        let password = &[0u8; 32];
        let padded = pad_password(password);
        assert_eq!(padded.len(), 32);
        assert_eq!(&padded[..], password);
    }

    #[test]
    fn test_constant_time_compare_equal() {
        let a = b"test1234test1234";
        let b = b"test1234test1234";
        assert!(constant_time_compare(a, b));
    }

    #[test]
    fn test_constant_time_compare_not_equal() {
        let a = b"test1234test1234";
        let b = b"test1234test1235";
        assert!(!constant_time_compare(a, b));
    }

    #[test]
    fn test_constant_time_compare_different_length() {
        let a = b"test";
        let b = b"testing";
        assert!(!constant_time_compare(a, b));
    }

    #[test]
    fn test_compute_encryption_key() {
        let password = b"user";
        let owner_key = &[0u8; 32];
        let permissions = -1;
        let file_id = b"test_file_id";
        let revision = 2;
        let key_length = 5;

        let key = compute_encryption_key(
            password,
            owner_key,
            permissions,
            file_id,
            revision,
            key_length,
            true,
        );

        assert_eq!(key.len(), key_length);
    }

    #[test]
    fn test_owner_password_hash_r2() {
        let owner = b"owner";
        let user = b"user";
        let revision = 2;
        let key_length = 5; // 40-bit

        let owner_hash = compute_owner_password_hash(owner, user, revision, key_length);

        // Should produce 32-byte hash
        assert_eq!(owner_hash.len(), 32);

        // Verify the hash can decrypt back to the user password
        // For R=2, decrypt with same RC4 key should give padded user password
    }

    #[test]
    fn test_owner_password_hash_r3() {
        let owner = b"owner";
        let user = b"user";
        let revision = 3;
        let key_length = 16; // 128-bit

        let owner_hash = compute_owner_password_hash(owner, user, revision, key_length);

        // Should produce 32-byte hash
        assert_eq!(owner_hash.len(), 32);
    }

    #[test]
    fn test_owner_password_hash_empty_owner() {
        // When owner password is empty, user password should be used
        let user = b"user";
        let revision = 3;
        let key_length = 16;

        let hash1 = compute_owner_password_hash(b"", user, revision, key_length);
        let hash2 = compute_owner_password_hash(user, user, revision, key_length);

        // Both should produce the same result since empty owner uses user password
        assert_eq!(hash1, hash2);
    }

    #[test]
    fn test_user_password_hash_r2() {
        let key = [0u8; 5]; // 40-bit key
        let file_id = b"test_file_id";
        let revision = 2;

        let user_hash = compute_user_password_hash(&key, file_id, revision);

        // R=2 always produces 32-byte result
        assert_eq!(user_hash.len(), 32);
    }

    #[test]
    fn test_user_password_hash_r3() {
        let key = [0u8; 16]; // 128-bit key
        let file_id = b"test_file_id";
        let revision = 3;

        let user_hash = compute_user_password_hash(&key, file_id, revision);

        // R>=3 produces 32-byte result (16 hash + 16 arbitrary)
        assert_eq!(user_hash.len(), 32);
    }

    #[test]
    fn test_encryption_roundtrip_r2() {
        // Test that we can create owner/user hashes and authenticate
        let owner_pass = b"owner123";
        let user_pass = b"user123";
        let file_id = b"test_file_id_123";
        let permissions = -1i32;
        let revision = 2;
        let key_length = 5;

        // Step 1: Compute owner hash (O value)
        let owner_hash = compute_owner_password_hash(owner_pass, user_pass, revision, key_length);

        // Step 2: Compute encryption key from user password
        let encryption_key = compute_encryption_key(
            user_pass,
            &owner_hash,
            permissions,
            file_id,
            revision,
            key_length,
            true,
        );

        // Step 3: Compute user hash (U value)
        let user_hash = compute_user_password_hash(&encryption_key, file_id, revision);

        // Step 4: Verify authentication works
        let auth_result = authenticate_user_password(
            user_pass,
            &user_hash,
            &owner_hash,
            permissions,
            file_id,
            revision,
            key_length,
            true,
        );

        assert!(auth_result.is_some());
        assert_eq!(auth_result.unwrap(), encryption_key);
    }

    #[test]
    fn test_encryption_roundtrip_r3() {
        // Test with R=3 (128-bit encryption)
        let owner_pass = b"owner456";
        let user_pass = b"user456";
        let file_id = b"test_file_id_456";
        let permissions = -1i32;
        let revision = 3;
        let key_length = 16;

        let owner_hash = compute_owner_password_hash(owner_pass, user_pass, revision, key_length);
        let encryption_key = compute_encryption_key(
            user_pass,
            &owner_hash,
            permissions,
            file_id,
            revision,
            key_length,
            true,
        );
        let user_hash = compute_user_password_hash(&encryption_key, file_id, revision);

        let auth_result = authenticate_user_password(
            user_pass,
            &user_hash,
            &owner_hash,
            permissions,
            file_id,
            revision,
            key_length,
            true,
        );

        assert!(auth_result.is_some());
        assert_eq!(auth_result.unwrap(), encryption_key);
    }
}
