//! Integration tests for the hybrid KEM, HKDF key derivation, and AEAD session.
//!
//! These exercise the public API exactly as a downstream crate (e.g. the Aura handshake) would.

use aura_crypto::{derive_session_keys, AeadSession, HybridPrivateKey, HybridSharedSecret};

/// keygen -> encapsulate -> decapsulate; both shared-secret halves must agree.
#[test]
fn test_hybrid_roundtrip() {
    let (private, public) = HybridPrivateKey::generate();

    let (ciphertext, ss_server) = public.encapsulate();
    let ss_client = private
        .decapsulate(&ciphertext)
        .expect("decapsulation succeeds");

    assert_eq!(
        ss_server.x25519_ss, ss_client.x25519_ss,
        "x25519 shared secret halves must match"
    );
    assert_eq!(
        ss_server.kyber_ss, ss_client.kyber_ss,
        "ML-KEM shared secret halves must match"
    );
    // Sanity: the two halves are independent 32-byte values.
    assert_eq!(ss_server.x25519_ss.len(), 32);
    assert_eq!(ss_server.kyber_ss.len(), 32);
}

/// Run the hybrid roundtrip many times to catch any rare encode/decode mismatch.
#[test]
fn test_hybrid_roundtrip_property() {
    for _ in 0..50 {
        let (private, public) = HybridPrivateKey::generate();
        let (ct, ss_server) = public.encapsulate();
        let ss_client = private.decapsulate(&ct).expect("decapsulation succeeds");
        assert_eq!(ss_server.x25519_ss, ss_client.x25519_ss);
        assert_eq!(ss_server.kyber_ss, ss_client.kyber_ss);
    }
}

/// Two independent keypairs must (overwhelmingly) not produce colliding ciphertexts/secrets, and
/// decapsulating someone else's ciphertext with the wrong key must NOT yield the server's secret.
#[test]
fn test_hybrid_wrong_key_disagrees() {
    let (private_a, _public_a) = HybridPrivateKey::generate();
    let (_private_b, public_b) = HybridPrivateKey::generate();

    // Encapsulate to B, then try to decapsulate with A's key.
    let (ct_for_b, ss_b) = public_b.encapsulate();
    let ss_a = private_a
        .decapsulate(&ct_for_b)
        .expect("decapsulation is infallible on a well-formed ciphertext");

    // x25519 half must differ (different static keys).
    assert_ne!(ss_a.x25519_ss, ss_b.x25519_ss);
    // ML-KEM half differs too: implicit rejection yields an unrelated pseudo-random secret.
    assert_ne!(ss_a.kyber_ss, ss_b.kyber_ss);
}

/// Helper: build a `HybridSharedSecret` from raw halves for KDF tests.
fn shared_from(x: [u8; 32], k: [u8; 32]) -> HybridSharedSecret {
    // We cannot construct `HybridSharedSecret` field-by-field from outside via a constructor, but
    // its fields are public, so build it directly.
    HybridSharedSecret {
        x25519_ss: x,
        kyber_ss: k.to_vec(),
    }
}

/// Same (shared, nonces) -> identical keys; changing a nonce -> different keys.
#[test]
fn test_kdf_deterministic() {
    let shared = shared_from([7u8; 32], [9u8; 32]);
    let client_nonce = [1u8; 32];
    let server_nonce = [2u8; 32];

    let k1 = derive_session_keys(&shared, &client_nonce, &server_nonce);
    let k2 = derive_session_keys(&shared, &client_nonce, &server_nonce);

    assert_eq!(k1.client_to_server, k2.client_to_server);
    assert_eq!(k1.server_to_client, k2.server_to_client);

    // The two directional keys are distinct (different halves of the HKDF output).
    assert_ne!(k1.client_to_server, k1.server_to_client);

    // Changing the client nonce changes the keys.
    let mut other_client = client_nonce;
    other_client[0] ^= 0xFF;
    let k3 = derive_session_keys(&shared, &other_client, &server_nonce);
    assert_ne!(k1.client_to_server, k3.client_to_server);
    assert_ne!(k1.server_to_client, k3.server_to_client);

    // Changing the server nonce changes the keys.
    let mut other_server = server_nonce;
    other_server[31] ^= 0x01;
    let k4 = derive_session_keys(&shared, &client_nonce, &other_server);
    assert_ne!(k1.client_to_server, k4.client_to_server);

    // Changing the shared secret changes the keys.
    let shared2 = shared_from([8u8; 32], [9u8; 32]);
    let k5 = derive_session_keys(&shared2, &client_nonce, &server_nonce);
    assert_ne!(k1.client_to_server, k5.client_to_server);
}

/// End-to-end: derive keys from a real handshake, then check the directional keys actually
/// protect traffic.
#[test]
fn test_kdf_from_real_handshake() {
    let (private, public) = HybridPrivateKey::generate();
    let (ct, ss_server) = public.encapsulate();
    let ss_client = private.decapsulate(&ct).expect("decapsulate");

    let client_nonce = [0x11u8; 32];
    let server_nonce = [0x22u8; 32];

    let server_keys = derive_session_keys(&ss_server, &client_nonce, &server_nonce);
    let client_keys = derive_session_keys(&ss_client, &client_nonce, &server_nonce);

    // Both sides derive the same key material.
    assert_eq!(server_keys.client_to_server, client_keys.client_to_server);
    assert_eq!(server_keys.server_to_client, client_keys.server_to_client);
}

/// seal then open returns the plaintext when AAD matches.
#[test]
fn test_aead_roundtrip() {
    let key = [0x42u8; 32];
    let mut sender = AeadSession::new(key);
    let mut receiver = AeadSession::new(key);

    let plaintext = b"hybrid post-quantum VPN payload";
    let aad = b"aura-header-v1";

    let ct = sender.seal(plaintext, aad);
    // Ciphertext is plaintext length + 16-byte Poly1305 tag.
    assert_eq!(ct.len(), plaintext.len() + 16);

    let recovered = receiver
        .open(&ct, aad)
        .expect("open succeeds with matching AAD");
    assert_eq!(recovered, plaintext);
}

/// Multiple sequential messages stay aligned between a sender and receiver session.
#[test]
fn test_aead_sequential_messages() {
    let key = [0x01u8; 32];
    let mut sender = AeadSession::new(key);
    let mut receiver = AeadSession::new(key);

    for i in 0u32..100 {
        let msg = format!("message number {i}");
        let aad = i.to_le_bytes();
        let ct = sender.seal(msg.as_bytes(), &aad);
        let pt = receiver.open(&ct, &aad).expect("aligned open succeeds");
        assert_eq!(pt, msg.as_bytes());
    }
}

/// Flipping a ciphertext byte, changing AAD, or using the wrong key must fail authentication.
#[test]
fn test_aead_tamper_detection() {
    let key = [0x42u8; 32];
    let aad = b"aura-header-v1";
    let plaintext = b"top secret";

    // 1. Flip a ciphertext byte.
    {
        let mut sender = AeadSession::new(key);
        let mut receiver = AeadSession::new(key);
        let mut ct = sender.seal(plaintext, aad);
        ct[0] ^= 0x01;
        assert!(
            receiver.open(&ct, aad).is_err(),
            "tampered ciphertext must fail"
        );
    }

    // 2. Flip a tag byte (last byte).
    {
        let mut sender = AeadSession::new(key);
        let mut receiver = AeadSession::new(key);
        let mut ct = sender.seal(plaintext, aad);
        let last = ct.len() - 1;
        ct[last] ^= 0x80;
        assert!(receiver.open(&ct, aad).is_err(), "tampered tag must fail");
    }

    // 3. Change the AAD.
    {
        let mut sender = AeadSession::new(key);
        let mut receiver = AeadSession::new(key);
        let ct = sender.seal(plaintext, aad);
        assert!(
            receiver.open(&ct, b"different-aad").is_err(),
            "mismatched AAD must fail"
        );
    }

    // 4. Wrong key.
    {
        let mut sender = AeadSession::new(key);
        let mut receiver = AeadSession::new([0x00u8; 32]);
        let ct = sender.seal(plaintext, aad);
        assert!(receiver.open(&ct, aad).is_err(), "wrong key must fail");
    }
}

/// A failed `open` still advances the counter, keeping a (sender, receiver) pair aligned for the
/// next message (so a single dropped/tampered frame does not desynchronize the stream here).
#[test]
fn test_aead_counter_advances_on_failure() {
    let key = [0x55u8; 32];
    let mut sender = AeadSession::new(key);
    let mut receiver = AeadSession::new(key);

    // Message 0: tamper -> fails, but receiver counter advances to 1.
    let mut ct0 = sender.seal(b"first", b"a");
    ct0[0] ^= 0x01;
    assert!(receiver.open(&ct0, b"a").is_err());

    // Message 1: both at counter 1 now -> succeeds.
    let ct1 = sender.seal(b"second", b"b");
    let pt1 = receiver.open(&ct1, b"b").expect("counters re-aligned at 1");
    assert_eq!(pt1, b"second");
}

/// 10_000 seal calls must use 10_000 distinct nonces. We verify this behaviorally: encrypting the
/// *same* plaintext+AAD under the same key 10_000 times yields 10_000 distinct ciphertexts, which
/// can only happen if the (counter-derived) nonce never repeats.
#[test]
fn test_nonce_no_repeat() {
    use std::collections::HashSet;

    let mut session = AeadSession::new([0x7Au8; 32]);
    let plaintext = b"constant";
    let aad = b"constant-aad";

    let mut seen: HashSet<Vec<u8>> = HashSet::with_capacity(10_000);
    for _ in 0..10_000 {
        let ct = session.seal(plaintext, aad);
        assert!(
            seen.insert(ct),
            "nonce reuse produced a duplicate ciphertext"
        );
    }
    assert_eq!(seen.len(), 10_000);
}