// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! Tests for secret sharing functionality
//!
//! This module contains integration tests and cross-module tests
//! for the secret sharing crate.

use x25519_dalek::{PublicKey as X25519PublicKey, StaticSecret};

use crate::{
    constants::{HPKE_CIPHERTEXT_SIZE, USER_SECRET_AAD},
    traits::SecretSharingEngine,
    types::{Aad, KeyId},
    wrapping::HpkeX25519ChaChaEngine,
    PublicKey, SecretSharingError, SecretSharingKey,
};

// Small helper to flip a single bit in a buffer.
fn flip(b: &mut [u8], i: usize) {
    b[i] ^= 0x01;
}

// This test simulates two parties (in our case, two TEEs):
//   • Sender: holds the global SecretSharingKey (SK) that needs to be
//     distributed to peers.
//   • Recipient: owns a wrapping X25519 keypair and should receive
//     SK_priv wrapped to its wrapping public key via HPKE.
//
// We check that:
//   1) HPKE wrapping (Base mode) + unwrapping round-trips SK_priv exactly.
//   2) The wire format for the ciphertext is 48 bytes:
//        ciphertext = body(32 bytes = SK_priv) || tag(16 bytes).
//   3) AAD is bound: unwrapping only succeeds when the exact AAD bytes match.
#[test]
fn hpke_roundtrip_wrap_unwrap() {
    // Recipient's static wrapping keypair
    let recipient_priv = StaticSecret::from([2u8; 32]);
    let recipient_pub = PublicKey::from(X25519PublicKey::from(&recipient_priv));
    let recipient_engine = HpkeX25519ChaChaEngine::with_wrapping_priv(recipient_priv);

    // Sender constructs its engine (with its own transport key; in HPKE Base mode
    // this key is *not* used for sender authentication, but we include the sender
    // pubkey in the envelope for provenance/telemetry). The SK generated here is
    // the *global* SecretSharingKey we intend to distribute to the recipient TEE.
    let sender_engine = HpkeX25519ChaChaEngine::with_wrapping_priv(StaticSecret::from([1u8; 32]));
    let sk = sender_engine
        .generate_secret_key()
        .expect("Failed to generate secret key");
    let aad = Aad {
        request_id: 42,
        key_id: KeyId::new(1).expect("Failed to create KeyId with value 1"),
    };

    // Sender wraps SK_priv to the recipient's transport public key.
    // Envelope fields:
    // - enc:        HPKE encapped key (32 bytes)
    // - ciphertext: 48 bytes = 32-byte body (SK_priv) || 16-byte AEAD tag
    // - metadata:   sender_wrapping_pub_x25519, recipient_wrapping_pub_x25519, key_id
    let envelope = sender_engine
        .wrap_for_recipient(&sk, recipient_pub, aad)
        .expect("Failed to wrap secret key for recipient");
    assert_eq!(envelope.ciphertext.len(), HPKE_CIPHERTEXT_SIZE);

    // Recipient unwraps in its enclave using its HPKE private key. This enforces:
    // - recipient match (WrongRecipient otherwise),
    // - correct decapsulation of `enc`,
    // - AEAD tag validity,
    // - exact AAD match.
    let sk2 = recipient_engine
        .unwrap_and_install(&envelope)
        .expect("Failed to unwrap and install secret key");

    // Check roundtrip
    // Both SK_priv and SK_pub must match exactly.
    assert_eq!(sk.public_key_bytes(), sk2.public_key_bytes());
    assert_eq!(
        sk.private_key_bytes_for_tests(),
        sk2.private_key_bytes_for_tests()
    );
}

/// Unwrap must fail if the envelope is presented to the wrong recipient.
#[test]
fn unwrap_rejects_wrong_recipient() {
    let good_recipient_priv = StaticSecret::from([3u8; 32]);
    let wrong_recipient_priv = StaticSecret::from([4u8; 32]);
    let good_recipient_pub = PublicKey::from(X25519PublicKey::from(&good_recipient_priv));

    let wrong_recipient_engine = HpkeX25519ChaChaEngine::with_wrapping_priv(wrong_recipient_priv);

    let sender_engine = HpkeX25519ChaChaEngine::with_wrapping_priv(StaticSecret::from([9u8; 32]));
    let sk = sender_engine
        .generate_secret_key()
        .expect("Failed to generate secret key for wrong recipient test");
    let aad = Aad {
        request_id: 1,
        key_id: KeyId::new(7).expect("Failed to create KeyId with value 7"),
    };

    let env = sender_engine
        .wrap_for_recipient(&sk, good_recipient_pub, aad)
        .expect("Failed to wrap secret key for wrong recipient test");

    let result = wrong_recipient_engine.unwrap_and_install(&env);
    assert!(matches!(result, Err(SecretSharingError::RecipientMismatch)));
}

/// AAD binding test: any change to AAD must cause open to fail.
#[test]
fn unwrap_rejects_mismatched_aad() {
    // Recipient's static wrapping keypair
    let recip_priv = StaticSecret::from([5u8; 32]);
    let recip_pub = PublicKey::from(X25519PublicKey::from(&recip_priv));
    let recipient_engine = HpkeX25519ChaChaEngine::with_wrapping_priv(recip_priv);

    // Sender engine + SK to distribute
    let sender_engine = HpkeX25519ChaChaEngine::with_wrapping_priv(StaticSecret::from([6u8; 32]));
    let sk = sender_engine
        .generate_secret_key()
        .expect("Failed to generate secret key for AAD mismatch test");

    // Wrap with a known AAD (the AAD values will be embedded/derived from the envelope)
    let aad_ok = Aad {
        request_id: 777,
        key_id: KeyId::new(42).expect("Failed to create KeyId with value 42"),
    };
    let envelope = sender_engine
        .wrap_for_recipient(&sk, recip_pub, aad_ok)
        .expect("Failed to wrap secret key for AAD mismatch test");
    assert_eq!(envelope.ciphertext.len(), HPKE_CIPHERTEXT_SIZE);

    // --- Tamper with the envelope AAD to simulate mismatch ---
    let mut bad_envelope = envelope.clone();

    // Change one AAD component carried in the envelope.
    // This ensures the receiver derives a different AAD than was used to seal.
    bad_envelope.aad.key_id = KeyId::new(envelope.aad.key_id.value() ^ 1)
        .expect("Failed to create modified KeyId for AAD tampering test"); // flip a bit / change value

    // Unwrap must fail with an AEAD/authentication error.
    let result = recipient_engine.unwrap_and_install(&bad_envelope);
    assert!(matches!(result, Err(SecretSharingError::UnwrapFailed(_))));

    // Test tampering with ciphertext tag
    let sender_engine = HpkeX25519ChaChaEngine::with_wrapping_priv(StaticSecret::from([8u8; 32]));
    let sk = sender_engine
        .generate_secret_key()
        .expect("Failed to generate secret key for tag tampering test");
    let aad = Aad {
        request_id: 9,
        key_id: KeyId::new(9).expect("Failed to create KeyId with value 9"),
    };

    let mut envelope = sender_engine
        .wrap_for_recipient(&sk, recip_pub, aad)
        .expect("Failed to wrap secret key for tag tampering test");

    // Flip last byte (inside the 16-byte tag suffix).
    flip(&mut envelope.ciphertext, 47);
    let result = recipient_engine.unwrap_and_install(&envelope);
    assert!(matches!(result, Err(SecretSharingError::UnwrapFailed(_))));

    // Test tampering with ciphertext body
    let sender_engine = HpkeX25519ChaChaEngine::with_wrapping_priv(StaticSecret::from([12u8; 32]));
    let sk = sender_engine
        .generate_secret_key()
        .expect("Failed to generate secret key for body tampering test");
    let aad = Aad {
        request_id: 10,
        key_id: KeyId::new(10).expect("Failed to create KeyId with value 10"),
    };

    let mut envelope = sender_engine
        .wrap_for_recipient(&sk, recip_pub, aad)
        .expect("Failed to wrap secret key for body tampering test");

    // Flip a byte in the 32-byte body (index 0..31).
    flip(&mut envelope.ciphertext, 0);
    let result = recipient_engine.unwrap_and_install(&envelope);
    assert!(matches!(result, Err(SecretSharingError::UnwrapFailed(_))));

    // Test tampering with HPKE encapsulated key
    let sender_engine = HpkeX25519ChaChaEngine::with_wrapping_priv(StaticSecret::from([14u8; 32]));
    let sk = sender_engine
        .generate_secret_key()
        .expect("Failed to generate secret key for encapsulated key tampering test");
    let aad = Aad {
        request_id: 11,
        key_id: KeyId::new(11).expect("Failed to create KeyId with value 11"),
    };

    let mut envelope = sender_engine
        .wrap_for_recipient(&sk, recip_pub, aad)
        .expect("Failed to wrap secret key for encapsulated key tampering test");

    // Break HPKE decapsulation.
    flip(&mut envelope.enc, 0);
    let result = recipient_engine.unwrap_and_install(&envelope);
    assert!(matches!(result, Err(SecretSharingError::UnwrapFailed(_))));
}

/// Test that the engine's wrapping public key can be retrieved and matches expectations.
#[test]
fn test_wrapping_public_key() {
    let priv_key = StaticSecret::from([15u8; 32]);
    let expected_pub = PublicKey::from(X25519PublicKey::from(&priv_key));

    let engine = HpkeX25519ChaChaEngine::with_wrapping_priv(priv_key);
    assert_eq!(engine.wrapping_public_key(), expected_pub);
}

/// Sanity: two independently generated SKs are (overwhelmingly likely to be) different.
/// (Keep as a soft signal of using a CSPRNG — not a strict randomness test.)
#[test]
fn generate_secret_key_is_randomish() {
    let engine = HpkeX25519ChaChaEngine::with_wrapping_priv(StaticSecret::from([31u8; 32]));
    let a: SecretSharingKey = engine
        .generate_secret_key()
        .expect("Failed to generate first secret key for randomness test");
    let b: SecretSharingKey = engine
        .generate_secret_key()
        .expect("Failed to generate second secret key for randomness test");
    assert_ne!(
        a.private_key_bytes_for_tests(),
        b.private_key_bytes_for_tests()
    );
}

/// Integration test for user payload encryption/decryption round-trip.
/// Tests the full cycle of client-side encryption to SecretSharingKey decryption.
#[test]
fn test_hpke_encrypt_user_payload_basic_roundtrip() {
    let sk = SecretSharingKey::from_test_seed(b"test-secret-key-payload");
    let sk_pub = sk.public_key_bytes();
    let payload = b"Hello, HPKE world!";
    let aad = USER_SECRET_AAD;

    // Encrypt the payload using the static method
    let ciphertext = HpkeX25519ChaChaEngine::hpke_encrypt_user_payload(&sk_pub, payload, aad)
        .expect("Encryption should succeed");

    // Verify ciphertext format: [enc (32 bytes)] || [payload + tag]
    assert!(
        ciphertext.len() >= 32 + payload.len() + 16,
        "Ciphertext should be at least enc + payload + tag"
    );

    // Split ciphertext
    let (enc_bytes, ct_with_tag) = ciphertext.split_at(32);
    let mut enc_array = [0u8; 32];
    enc_array.copy_from_slice(enc_bytes);

    // Decrypt using the static function
    let sk_priv = sk.private_key().to_bytes();
    let decrypted =
        HpkeX25519ChaChaEngine::hpke_decrypt_user_payload(&sk_priv, &enc_array, ct_with_tag, aad)
            .expect("Decryption should succeed");

    assert_eq!(
        &payload[..],
        &*decrypted,
        "Roundtrip should preserve payload"
    );
}

/// Integration test for AAD binding in user payload encryption.
/// Verifies that different AAD values produce different ciphertexts and prevent wrong-AAD decryption.
#[test]
fn test_hpke_encrypt_user_payload_different_aad() {
    let sk = SecretSharingKey::from_test_seed(b"test-secret-key-payload-aad");
    let sk_pub = sk.public_key_bytes();
    let sk_priv = sk.private_key().to_bytes();
    let payload = b"test payload for AAD testing";

    let aad1 = b"aad_value_1";
    let aad2 = b"aad_value_2";

    // Encrypt with first AAD
    let ciphertext1 = HpkeX25519ChaChaEngine::hpke_encrypt_user_payload(&sk_pub, payload, aad1)
        .expect("Encryption should succeed");

    // Encrypt with second AAD
    let ciphertext2 = HpkeX25519ChaChaEngine::hpke_encrypt_user_payload(&sk_pub, payload, aad2)
        .expect("Encryption should succeed");

    // Ciphertexts should be different due to different AAD
    assert_ne!(
        ciphertext1, ciphertext2,
        "Different AAD should produce different ciphertexts"
    );

    // Decrypt with correct AAD should work
    let (enc1, ct1) = ciphertext1.split_at(32);
    let mut enc1_array = [0u8; 32];
    enc1_array.copy_from_slice(enc1);

    let decrypted1 =
        HpkeX25519ChaChaEngine::hpke_decrypt_user_payload(&sk_priv, &enc1_array, ct1, aad1)
            .expect("Decryption with correct AAD should succeed");

    assert_eq!(payload, &**decrypted1);

    // Decrypt with wrong AAD should fail
    let result =
        HpkeX25519ChaChaEngine::hpke_decrypt_user_payload(&sk_priv, &enc1_array, ct1, aad2);
    assert!(result.is_err(), "Decryption with wrong AAD should fail");
}

/// Integration test for user payload encryption format validation.
/// Verifies the ciphertext format and structure conform to specification.
#[test]
fn test_hpke_encrypt_user_payload_format_validation() {
    let sk = SecretSharingKey::from_test_seed(b"test-secret-key-format");
    let sk_pub = sk.public_key_bytes();
    let payload = b"format validation test";
    let aad = USER_SECRET_AAD;

    let ciphertext = HpkeX25519ChaChaEngine::hpke_encrypt_user_payload(&sk_pub, payload, aad)
        .expect("Encryption should succeed");

    // Verify expected format: [enc (32 bytes)] || [ciphertext + tag]
    assert!(
        ciphertext.len() >= 48,
        "Minimum ciphertext length should be 48 bytes"
    );
    assert_eq!(
        ciphertext.len(),
        32 + payload.len() + 16,
        "Total length should be enc + payload + tag"
    );

    // First 32 bytes should be the encapsulated key
    let enc_part = &ciphertext[0..32];
    assert!(
        enc_part.iter().any(|&b| b != 0),
        "Encapsulated key should not be all zeros"
    );

    // Remaining part should be ciphertext + tag
    let ct_with_tag = &ciphertext[32..];
    assert_eq!(
        ct_with_tag.len(),
        payload.len() + 16,
        "Ciphertext part should be payload + 16-byte tag"
    );
}

/// Concurrent access test for HPKE encryption.
/// Verifies that multiple threads can safely encrypt payloads simultaneously.
#[test]
fn test_concurrent_hpke_encryption() {
    use std::{sync::Arc, thread};

    let sk = SecretSharingKey::from_test_seed(b"concurrent-test-key");
    let sk_pub = Arc::new(sk.public_key_bytes());
    let payload = b"concurrent test payload";

    let handles: Vec<_> = (0..10)
        .map(|_i| {
            let sk_pub = Arc::clone(&sk_pub);
            thread::spawn(move || {
                HpkeX25519ChaChaEngine::hpke_encrypt_user_payload(&sk_pub, payload, USER_SECRET_AAD)
            })
        })
        .collect();

    for handle in handles {
        handle
            .join()
            .unwrap()
            .expect("Concurrent encryption should succeed");
    }
}

/// Test HPKE decrypt with various malformed ciphertext inputs.
/// Verifies robust error handling for invalid/corrupted data.
#[test]
fn test_hpke_decrypt_malformed_ciphertext() {
    use assert_matches::assert_matches;

    let sk = SecretSharingKey::from_test_seed(b"malformed-test-key");
    let sk_priv = sk.private_key().to_bytes();
    let enc = [0u8; 32];

    // Test various malformed inputs
    let test_cases = [
        vec![],           // Empty
        vec![0u8; 15],    // Too short for tag
        vec![0u8; 16],    // Only tag, no body
        vec![0xff; 1000], // Invalid data
    ];

    for malformed_ct in test_cases {
        let result = HpkeX25519ChaChaEngine::hpke_decrypt_user_payload(
            &sk_priv,
            &enc,
            &malformed_ct,
            USER_SECRET_AAD,
        );
        assert_matches!(
            result,
            Err(SecretSharingError::EnvelopeInvalid(_)) | Err(SecretSharingError::UnwrapFailed(_)),
            "Malformed ciphertext should fail with appropriate error type"
        );
    }
}

/// Test HPKE encryption/decryption with zero-length payload.
/// Verifies that empty payloads can be encrypted and decrypted correctly.
#[test]
fn test_hpke_encrypt_zero_length_payload() {
    let sk = SecretSharingKey::from_test_seed(b"zero-length-test-key");
    let sk_pub = sk.public_key_bytes();
    let sk_priv = sk.private_key().to_bytes();
    let empty_payload = b""; // Zero-length payload
    let aad = USER_SECRET_AAD;

    // Encrypt the empty payload
    let ciphertext = HpkeX25519ChaChaEngine::hpke_encrypt_user_payload(&sk_pub, empty_payload, aad)
        .expect("Encryption of zero-length payload should succeed");

    // Verify ciphertext format: [enc (32 bytes)] || [tag (16 bytes)] = 48 bytes for empty payload
    assert_eq!(
        ciphertext.len(),
        48,
        "Ciphertext for zero-length payload should be exactly 48 bytes (32 enc + 0 payload + 16 tag)"
    );

    // Split ciphertext: enc (32 bytes) + tag (16 bytes)
    let (enc_bytes, ct_with_tag) = ciphertext.split_at(32);
    let mut enc_array = [0u8; 32];
    enc_array.copy_from_slice(enc_bytes);

    // Verify the remaining part is just the tag (16 bytes)
    assert_eq!(
        ct_with_tag.len(),
        16,
        "For zero-length payload, ciphertext part should be only the 16-byte tag"
    );

    // Decrypt the empty payload
    let decrypted =
        HpkeX25519ChaChaEngine::hpke_decrypt_user_payload(&sk_priv, &enc_array, ct_with_tag, aad)
            .expect("Decryption of zero-length payload should succeed");

    // Verify round-trip: decrypted should be empty
    assert_eq!(
        decrypted.len(),
        0,
        "Decrypted zero-length payload should be empty"
    );
    assert_eq!(
        &**decrypted, empty_payload,
        "Round-trip should preserve zero-length payload"
    );
}
