//! TLS 1.3 key schedule (RFC 8446 §7).
//!
//! Derives encryption keys, IVs, and Finished verify data from shared secrets
//! using the HKDF-based key schedule defined in RFC 8446.

use we_crypto::hkdf::{hkdf_expand, hkdf_extract};
use we_crypto::hmac::{HashFunction, Hmac};
use we_crypto::sha2::{Sha256, Sha384};

use super::record::CipherSuite;

// ---------------------------------------------------------------------------
// Hash abstraction for cipher suites
// ---------------------------------------------------------------------------

/// Trait linking a cipher suite to its hash function for the key schedule.
trait KeyScheduleHash: HashFunction + Clone {
    fn hash(data: &[u8]) -> Vec<u8>;
    fn hash_len() -> usize;
}

impl KeyScheduleHash for Sha256 {
    fn hash(data: &[u8]) -> Vec<u8> {
        we_crypto::sha2::sha256(data).to_vec()
    }
    fn hash_len() -> usize {
        32
    }
}

impl KeyScheduleHash for Sha384 {
    fn hash(data: &[u8]) -> Vec<u8> {
        we_crypto::sha2::sha384(data).to_vec()
    }
    fn hash_len() -> usize {
        48
    }
}

// ---------------------------------------------------------------------------
// HKDF-Expand-Label (RFC 8446 §7.1)
// ---------------------------------------------------------------------------

/// HKDF-Expand-Label as defined in RFC 8446 §7.1:
///
/// ```text
/// HKDF-Expand-Label(Secret, Label, Context, Length) =
///     HKDF-Expand(Secret, HkdfLabel, Length)
///
/// struct {
///     uint16 length = Length;
///     opaque label<7..255> = "tls13 " + Label;
///     opaque context<0..255> = Context;
/// } HkdfLabel;
/// ```
fn hkdf_expand_label<H: HashFunction>(
    secret: &[u8],
    label: &[u8],
    context: &[u8],
    length: usize,
) -> Vec<u8> {
    let full_label = [b"tls13 ", label].concat();

    // Build HkdfLabel structure
    let mut hkdf_label = Vec::with_capacity(2 + 1 + full_label.len() + 1 + context.len());
    hkdf_label.extend_from_slice(&(length as u16).to_be_bytes());
    hkdf_label.push(full_label.len() as u8);
    hkdf_label.extend_from_slice(&full_label);
    hkdf_label.push(context.len() as u8);
    hkdf_label.extend_from_slice(context);

    hkdf_expand::<H>(secret, &hkdf_label, length).expect("HKDF-Expand-Label length within bounds")
}

/// Derive-Secret as defined in RFC 8446 §7.1:
///
/// ```text
/// Derive-Secret(Secret, Label, Messages) =
///     HKDF-Expand-Label(Secret, Label, Transcript-Hash(Messages), Hash.length)
/// ```
fn derive_secret<H: KeyScheduleHash>(
    secret: &[u8],
    label: &[u8],
    transcript_hash: &[u8],
) -> Vec<u8> {
    hkdf_expand_label::<H>(secret, label, transcript_hash, H::hash_len())
}

// ---------------------------------------------------------------------------
// Transcript hash
// ---------------------------------------------------------------------------

/// Running transcript hash over handshake messages.
#[derive(Clone)]
pub struct TranscriptHash {
    suite: CipherSuite,
    sha256: Option<Sha256>,
    sha384: Option<Sha384>,
}

impl TranscriptHash {
    /// Create a new transcript hash for the given cipher suite.
    pub fn new(suite: CipherSuite) -> Self {
        match suite {
            CipherSuite::Aes128Gcm | CipherSuite::Chacha20Poly1305 => Self {
                suite,
                sha256: Some(Sha256::new()),
                sha384: None,
            },
            CipherSuite::Aes256Gcm => Self {
                suite,
                sha256: None,
                sha384: Some(Sha384::new()),
            },
        }
    }

    /// Update the transcript with a handshake message.
    pub fn update(&mut self, data: &[u8]) {
        if let Some(h) = &mut self.sha256 {
            h.update(data);
        }
        if let Some(h) = &mut self.sha384 {
            h.update(data);
        }
    }

    /// Get the current transcript hash value (does not consume the hasher).
    pub fn current_hash(&self) -> Vec<u8> {
        match self.suite {
            CipherSuite::Aes128Gcm | CipherSuite::Chacha20Poly1305 => {
                self.sha256.clone().unwrap().finalize().to_vec()
            }
            CipherSuite::Aes256Gcm => self.sha384.clone().unwrap().finalize().to_vec(),
        }
    }

    /// Hash length in bytes.
    pub fn hash_len(&self) -> usize {
        match self.suite {
            CipherSuite::Aes128Gcm | CipherSuite::Chacha20Poly1305 => 32,
            CipherSuite::Aes256Gcm => 48,
        }
    }
}

// ---------------------------------------------------------------------------
// Traffic keys
// ---------------------------------------------------------------------------

/// Derived traffic encryption key and IV.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TrafficKeys {
    pub key: Vec<u8>,
    pub iv: Vec<u8>,
}

// ---------------------------------------------------------------------------
// Key schedule (RFC 8446 §7.1)
// ---------------------------------------------------------------------------

/// TLS 1.3 key schedule state machine.
///
/// Progresses through: Early -> Handshake -> Master, deriving traffic secrets
/// at each stage.
pub struct KeySchedule {
    suite: CipherSuite,
    early_secret: Vec<u8>,
    handshake_secret: Option<Vec<u8>>,
    master_secret: Option<Vec<u8>>,
    client_handshake_traffic_secret: Option<Vec<u8>>,
    server_handshake_traffic_secret: Option<Vec<u8>>,
    client_app_traffic_secret: Option<Vec<u8>>,
    server_app_traffic_secret: Option<Vec<u8>>,
}

impl KeySchedule {
    /// Create a new key schedule, computing the Early Secret.
    ///
    /// For TLS 1.3 without PSK, pass `None` for `psk` (uses a zero-filled key).
    pub fn new(suite: CipherSuite, psk: Option<&[u8]>) -> Self {
        let hash_len = hash_len_for_suite(suite);
        let zero_key = vec![0u8; hash_len];
        let psk = psk.unwrap_or(&zero_key);
        let salt = vec![0u8; hash_len];

        let early_secret = extract_for_suite(suite, &salt, psk);

        Self {
            suite,
            early_secret,
            handshake_secret: None,
            master_secret: None,
            client_handshake_traffic_secret: None,
            server_handshake_traffic_secret: None,
            client_app_traffic_secret: None,
            server_app_traffic_secret: None,
        }
    }

    /// Derive handshake secrets from the shared ECDHE secret and transcript hash.
    ///
    /// The transcript hash should cover ClientHello...ServerHello.
    pub fn derive_handshake_secrets(&mut self, shared_secret: &[u8], transcript_hash: &[u8]) {
        let derived = self.derive_secret_inner(b"derived", &empty_hash(self.suite));
        let handshake_secret = extract_for_suite(self.suite, &derived, shared_secret);

        self.client_handshake_traffic_secret =
            Some(self.derive_secret_from(&handshake_secret, b"c hs traffic", transcript_hash));
        self.server_handshake_traffic_secret =
            Some(self.derive_secret_from(&handshake_secret, b"s hs traffic", transcript_hash));

        self.handshake_secret = Some(handshake_secret);
    }

    /// Derive application traffic secrets from the transcript hash.
    ///
    /// The transcript hash should cover ClientHello...server Finished.
    pub fn derive_app_secrets(&mut self, transcript_hash: &[u8]) {
        let hs_secret = self
            .handshake_secret
            .as_ref()
            .expect("must call derive_handshake_secrets first");

        let derived =
            derive_secret_for_suite(self.suite, hs_secret, b"derived", &empty_hash(self.suite));
        let zero_ikm = vec![0u8; hash_len_for_suite(self.suite)];
        let master_secret = extract_for_suite(self.suite, &derived, &zero_ikm);

        self.client_app_traffic_secret = Some(derive_secret_for_suite(
            self.suite,
            &master_secret,
            b"c ap traffic",
            transcript_hash,
        ));
        self.server_app_traffic_secret = Some(derive_secret_for_suite(
            self.suite,
            &master_secret,
            b"s ap traffic",
            transcript_hash,
        ));

        self.master_secret = Some(master_secret);
    }

    /// Get the client handshake traffic keys.
    pub fn client_handshake_keys(&self) -> TrafficKeys {
        let secret = self
            .client_handshake_traffic_secret
            .as_ref()
            .expect("handshake secrets not derived");
        self.derive_traffic_keys(secret)
    }

    /// Get the server handshake traffic keys.
    pub fn server_handshake_keys(&self) -> TrafficKeys {
        let secret = self
            .server_handshake_traffic_secret
            .as_ref()
            .expect("handshake secrets not derived");
        self.derive_traffic_keys(secret)
    }

    /// Get the client application traffic keys.
    pub fn client_app_keys(&self) -> TrafficKeys {
        let secret = self
            .client_app_traffic_secret
            .as_ref()
            .expect("app secrets not derived");
        self.derive_traffic_keys(secret)
    }

    /// Get the server application traffic keys.
    pub fn server_app_keys(&self) -> TrafficKeys {
        let secret = self
            .server_app_traffic_secret
            .as_ref()
            .expect("app secrets not derived");
        self.derive_traffic_keys(secret)
    }

    /// Compute the Finished verify data (RFC 8446 §4.4.4).
    ///
    /// `base_key` is the handshake traffic secret for the sender.
    /// `transcript_hash` is the hash of all handshake messages up to (but not
    /// including) the Finished message.
    pub fn compute_finished_verify_data(&self, base_key: &[u8], transcript_hash: &[u8]) -> Vec<u8> {
        let hash_len = hash_len_for_suite(self.suite);
        let finished_key = expand_label_for_suite(self.suite, base_key, b"finished", &[], hash_len);

        // HMAC(finished_key, transcript_hash)
        hmac_for_suite(self.suite, &finished_key, transcript_hash)
    }

    /// Get the client handshake traffic secret.
    pub fn client_handshake_traffic_secret(&self) -> Option<&[u8]> {
        self.client_handshake_traffic_secret.as_deref()
    }

    /// Get the server handshake traffic secret.
    pub fn server_handshake_traffic_secret(&self) -> Option<&[u8]> {
        self.server_handshake_traffic_secret.as_deref()
    }

    /// Get the client application traffic secret.
    pub fn client_app_traffic_secret(&self) -> Option<&[u8]> {
        self.client_app_traffic_secret.as_deref()
    }

    /// Get the server application traffic secret.
    pub fn server_app_traffic_secret(&self) -> Option<&[u8]> {
        self.server_app_traffic_secret.as_deref()
    }

    /// Get the early secret.
    pub fn early_secret(&self) -> &[u8] {
        &self.early_secret
    }

    /// Get the handshake secret.
    pub fn handshake_secret(&self) -> Option<&[u8]> {
        self.handshake_secret.as_deref()
    }

    /// Get the master secret.
    pub fn master_secret(&self) -> Option<&[u8]> {
        self.master_secret.as_deref()
    }

    /// Get the cipher suite.
    pub fn cipher_suite(&self) -> CipherSuite {
        self.suite
    }

    // -- internal helpers --

    fn derive_secret_inner(&self, label: &[u8], transcript_hash: &[u8]) -> Vec<u8> {
        derive_secret_for_suite(self.suite, &self.early_secret, label, transcript_hash)
    }

    fn derive_secret_from(&self, secret: &[u8], label: &[u8], transcript_hash: &[u8]) -> Vec<u8> {
        derive_secret_for_suite(self.suite, secret, label, transcript_hash)
    }

    fn derive_traffic_keys(&self, secret: &[u8]) -> TrafficKeys {
        let key_len = self.suite.key_len();
        let iv_len = self.suite.iv_len();
        let key = expand_label_for_suite(self.suite, secret, b"key", &[], key_len);
        let iv = expand_label_for_suite(self.suite, secret, b"iv", &[], iv_len);
        TrafficKeys { key, iv }
    }
}

// ---------------------------------------------------------------------------
// Suite-dispatched helpers
// ---------------------------------------------------------------------------

fn hash_len_for_suite(suite: CipherSuite) -> usize {
    match suite {
        CipherSuite::Aes128Gcm | CipherSuite::Chacha20Poly1305 => 32,
        CipherSuite::Aes256Gcm => 48,
    }
}

fn empty_hash(suite: CipherSuite) -> Vec<u8> {
    match suite {
        CipherSuite::Aes128Gcm | CipherSuite::Chacha20Poly1305 => Sha256::hash(&[]),
        CipherSuite::Aes256Gcm => Sha384::hash(&[]),
    }
}

fn extract_for_suite(suite: CipherSuite, salt: &[u8], ikm: &[u8]) -> Vec<u8> {
    match suite {
        CipherSuite::Aes128Gcm | CipherSuite::Chacha20Poly1305 => hkdf_extract::<Sha256>(salt, ikm),
        CipherSuite::Aes256Gcm => hkdf_extract::<Sha384>(salt, ikm),
    }
}

fn expand_label_for_suite(
    suite: CipherSuite,
    secret: &[u8],
    label: &[u8],
    context: &[u8],
    length: usize,
) -> Vec<u8> {
    match suite {
        CipherSuite::Aes128Gcm | CipherSuite::Chacha20Poly1305 => {
            hkdf_expand_label::<Sha256>(secret, label, context, length)
        }
        CipherSuite::Aes256Gcm => hkdf_expand_label::<Sha384>(secret, label, context, length),
    }
}

fn derive_secret_for_suite(
    suite: CipherSuite,
    secret: &[u8],
    label: &[u8],
    transcript_hash: &[u8],
) -> Vec<u8> {
    match suite {
        CipherSuite::Aes128Gcm | CipherSuite::Chacha20Poly1305 => {
            derive_secret::<Sha256>(secret, label, transcript_hash)
        }
        CipherSuite::Aes256Gcm => derive_secret::<Sha384>(secret, label, transcript_hash),
    }
}

fn hmac_for_suite(suite: CipherSuite, key: &[u8], data: &[u8]) -> Vec<u8> {
    match suite {
        CipherSuite::Aes128Gcm | CipherSuite::Chacha20Poly1305 => {
            let mut hmac = Hmac::<Sha256>::new(key);
            hmac.update(data);
            hmac.finalize()
        }
        CipherSuite::Aes256Gcm => {
            let mut hmac = Hmac::<Sha384>::new(key);
            hmac.update(data);
            hmac.finalize()
        }
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    // -- HKDF-Expand-Label tests --

    #[test]
    fn hkdf_expand_label_produces_correct_length() {
        let secret = vec![0x42u8; 32];
        let result = hkdf_expand_label::<Sha256>(&secret, b"key", &[], 16);
        assert_eq!(result.len(), 16);
    }

    #[test]
    fn hkdf_expand_label_produces_correct_length_48() {
        let secret = vec![0x42u8; 48];
        let result = hkdf_expand_label::<Sha384>(&secret, b"key", &[], 32);
        assert_eq!(result.len(), 32);
    }

    #[test]
    fn hkdf_expand_label_with_context() {
        let secret = vec![0x42u8; 32];
        let ctx = b"some context data";
        let result = hkdf_expand_label::<Sha256>(&secret, b"iv", ctx, 12);
        assert_eq!(result.len(), 12);
    }

    #[test]
    fn hkdf_expand_label_different_labels_different_output() {
        let secret = vec![0x42u8; 32];
        let r1 = hkdf_expand_label::<Sha256>(&secret, b"key", &[], 16);
        let r2 = hkdf_expand_label::<Sha256>(&secret, b"iv", &[], 16);
        assert_ne!(r1, r2);
    }

    #[test]
    fn hkdf_expand_label_different_contexts_different_output() {
        let secret = vec![0x42u8; 32];
        let r1 = hkdf_expand_label::<Sha256>(&secret, b"key", b"ctx1", 16);
        let r2 = hkdf_expand_label::<Sha256>(&secret, b"key", b"ctx2", 16);
        assert_ne!(r1, r2);
    }

    // -- Transcript hash tests --

    #[test]
    fn transcript_hash_sha256_empty() {
        let th = TranscriptHash::new(CipherSuite::Aes128Gcm);
        let hash = th.current_hash();
        assert_eq!(hash.len(), 32);
        // SHA-256 of empty string
        let expected = we_crypto::sha2::sha256(&[]);
        assert_eq!(hash, expected.to_vec());
    }

    #[test]
    fn transcript_hash_sha384_empty() {
        let th = TranscriptHash::new(CipherSuite::Aes256Gcm);
        let hash = th.current_hash();
        assert_eq!(hash.len(), 48);
        let expected = we_crypto::sha2::sha384(&[]);
        assert_eq!(hash, expected.to_vec());
    }

    #[test]
    fn transcript_hash_incremental() {
        let mut th = TranscriptHash::new(CipherSuite::Aes128Gcm);
        th.update(b"hello");
        th.update(b" world");
        let hash = th.current_hash();

        let expected = we_crypto::sha2::sha256(b"hello world");
        assert_eq!(hash, expected.to_vec());
    }

    #[test]
    fn transcript_hash_chacha_uses_sha256() {
        let th = TranscriptHash::new(CipherSuite::Chacha20Poly1305);
        assert_eq!(th.hash_len(), 32);
        let hash = th.current_hash();
        assert_eq!(hash.len(), 32);
    }

    #[test]
    fn transcript_hash_clone_independence() {
        let mut th = TranscriptHash::new(CipherSuite::Aes128Gcm);
        th.update(b"first");
        let th_clone = th.clone();
        th.update(b"second");

        let hash1 = th.current_hash();
        let hash2 = th_clone.current_hash();
        assert_ne!(hash1, hash2);
    }

    // -- Key schedule tests --

    #[test]
    fn key_schedule_early_secret_no_psk() {
        let ks = KeySchedule::new(CipherSuite::Aes128Gcm, None);
        assert_eq!(ks.early_secret().len(), 32);
    }

    #[test]
    fn key_schedule_early_secret_with_psk() {
        let psk = vec![0xABu8; 32];
        let ks1 = KeySchedule::new(CipherSuite::Aes128Gcm, Some(&psk));
        let ks2 = KeySchedule::new(CipherSuite::Aes128Gcm, None);
        // Different PSK should produce different early secret
        assert_ne!(ks1.early_secret(), ks2.early_secret());
    }

    #[test]
    fn key_schedule_early_secret_aes256() {
        let ks = KeySchedule::new(CipherSuite::Aes256Gcm, None);
        assert_eq!(ks.early_secret().len(), 48);
    }

    #[test]
    fn key_schedule_early_secret_chacha() {
        let ks = KeySchedule::new(CipherSuite::Chacha20Poly1305, None);
        assert_eq!(ks.early_secret().len(), 32);
    }

    #[test]
    fn key_schedule_derive_handshake_secrets() {
        let mut ks = KeySchedule::new(CipherSuite::Aes128Gcm, None);
        let shared_secret = vec![0x55u8; 32];
        let transcript = we_crypto::sha2::sha256(b"ClientHello...ServerHello");

        ks.derive_handshake_secrets(&shared_secret, &transcript);

        assert!(ks.handshake_secret().is_some());
        assert!(ks.client_handshake_traffic_secret().is_some());
        assert!(ks.server_handshake_traffic_secret().is_some());
        assert_eq!(ks.client_handshake_traffic_secret().unwrap().len(), 32);
        assert_eq!(ks.server_handshake_traffic_secret().unwrap().len(), 32);
        // Client and server traffic secrets must differ
        assert_ne!(
            ks.client_handshake_traffic_secret(),
            ks.server_handshake_traffic_secret()
        );
    }

    #[test]
    fn key_schedule_derive_app_secrets() {
        let mut ks = KeySchedule::new(CipherSuite::Aes128Gcm, None);
        let shared_secret = vec![0x55u8; 32];
        let hs_hash = we_crypto::sha2::sha256(b"handshake transcript");
        let app_hash = we_crypto::sha2::sha256(b"full transcript including Finished");

        ks.derive_handshake_secrets(&shared_secret, &hs_hash);
        ks.derive_app_secrets(&app_hash);

        assert!(ks.master_secret().is_some());
        assert!(ks.client_app_traffic_secret().is_some());
        assert!(ks.server_app_traffic_secret().is_some());
        assert_ne!(
            ks.client_app_traffic_secret(),
            ks.server_app_traffic_secret()
        );
    }

    #[test]
    fn key_schedule_traffic_keys_correct_lengths_aes128() {
        let mut ks = KeySchedule::new(CipherSuite::Aes128Gcm, None);
        let shared_secret = vec![0x55u8; 32];
        let transcript = we_crypto::sha2::sha256(b"test");

        ks.derive_handshake_secrets(&shared_secret, &transcript);

        let keys = ks.client_handshake_keys();
        assert_eq!(keys.key.len(), 16); // AES-128 key
        assert_eq!(keys.iv.len(), 12); // 12-byte IV
    }

    #[test]
    fn key_schedule_traffic_keys_correct_lengths_aes256() {
        let mut ks = KeySchedule::new(CipherSuite::Aes256Gcm, None);
        let shared_secret = vec![0x55u8; 48];
        let transcript = we_crypto::sha2::sha384(b"test");

        ks.derive_handshake_secrets(&shared_secret, &transcript);

        let keys = ks.client_handshake_keys();
        assert_eq!(keys.key.len(), 32); // AES-256 key
        assert_eq!(keys.iv.len(), 12);
    }

    #[test]
    fn key_schedule_traffic_keys_correct_lengths_chacha() {
        let mut ks = KeySchedule::new(CipherSuite::Chacha20Poly1305, None);
        let shared_secret = vec![0x55u8; 32];
        let transcript = we_crypto::sha2::sha256(b"test");

        ks.derive_handshake_secrets(&shared_secret, &transcript);

        let keys = ks.server_handshake_keys();
        assert_eq!(keys.key.len(), 32); // ChaCha20 key
        assert_eq!(keys.iv.len(), 12);
    }

    #[test]
    fn key_schedule_client_server_keys_differ() {
        let mut ks = KeySchedule::new(CipherSuite::Aes128Gcm, None);
        let shared_secret = vec![0x55u8; 32];
        let transcript = we_crypto::sha2::sha256(b"test");

        ks.derive_handshake_secrets(&shared_secret, &transcript);

        let client = ks.client_handshake_keys();
        let server = ks.server_handshake_keys();
        assert_ne!(client.key, server.key);
        assert_ne!(client.iv, server.iv);
    }

    #[test]
    fn key_schedule_app_keys_differ_from_handshake_keys() {
        let mut ks = KeySchedule::new(CipherSuite::Aes128Gcm, None);
        let shared_secret = vec![0x55u8; 32];
        let hs_hash = we_crypto::sha2::sha256(b"hs");
        let app_hash = we_crypto::sha2::sha256(b"app");

        ks.derive_handshake_secrets(&shared_secret, &hs_hash);
        let hs_keys = ks.client_handshake_keys();

        ks.derive_app_secrets(&app_hash);
        let app_keys = ks.client_app_keys();

        assert_ne!(hs_keys.key, app_keys.key);
        assert_ne!(hs_keys.iv, app_keys.iv);
    }

    #[test]
    fn key_schedule_different_shared_secrets_produce_different_keys() {
        let transcript = we_crypto::sha2::sha256(b"test");

        let mut ks1 = KeySchedule::new(CipherSuite::Aes128Gcm, None);
        ks1.derive_handshake_secrets(&vec![0x11u8; 32], &transcript);

        let mut ks2 = KeySchedule::new(CipherSuite::Aes128Gcm, None);
        ks2.derive_handshake_secrets(&vec![0x22u8; 32], &transcript);

        assert_ne!(
            ks1.client_handshake_keys().key,
            ks2.client_handshake_keys().key
        );
    }

    #[test]
    fn key_schedule_different_transcripts_produce_different_keys() {
        let shared_secret = vec![0x55u8; 32];

        let mut ks1 = KeySchedule::new(CipherSuite::Aes128Gcm, None);
        ks1.derive_handshake_secrets(&shared_secret, &we_crypto::sha2::sha256(b"transcript1"));

        let mut ks2 = KeySchedule::new(CipherSuite::Aes128Gcm, None);
        ks2.derive_handshake_secrets(&shared_secret, &we_crypto::sha2::sha256(b"transcript2"));

        assert_ne!(
            ks1.client_handshake_keys().key,
            ks2.client_handshake_keys().key
        );
    }

    // -- Finished verify data tests --

    #[test]
    fn finished_verify_data_correct_length() {
        let mut ks = KeySchedule::new(CipherSuite::Aes128Gcm, None);
        let shared_secret = vec![0x55u8; 32];
        let transcript = we_crypto::sha2::sha256(b"test");
        ks.derive_handshake_secrets(&shared_secret, &transcript);

        let verify_data = ks.compute_finished_verify_data(
            ks.client_handshake_traffic_secret().unwrap(),
            &transcript,
        );
        assert_eq!(verify_data.len(), 32);
    }

    #[test]
    fn finished_verify_data_correct_length_384() {
        let mut ks = KeySchedule::new(CipherSuite::Aes256Gcm, None);
        let shared_secret = vec![0x55u8; 48];
        let transcript = we_crypto::sha2::sha384(b"test");
        ks.derive_handshake_secrets(&shared_secret, &transcript);

        let verify_data = ks.compute_finished_verify_data(
            ks.client_handshake_traffic_secret().unwrap(),
            &transcript,
        );
        assert_eq!(verify_data.len(), 48);
    }

    #[test]
    fn finished_verify_data_deterministic() {
        let mut ks = KeySchedule::new(CipherSuite::Aes128Gcm, None);
        let shared_secret = vec![0x55u8; 32];
        let transcript = we_crypto::sha2::sha256(b"test");
        ks.derive_handshake_secrets(&shared_secret, &transcript);

        let secret = ks.client_handshake_traffic_secret().unwrap().to_vec();
        let vd1 = ks.compute_finished_verify_data(&secret, &transcript);
        let vd2 = ks.compute_finished_verify_data(&secret, &transcript);
        assert_eq!(vd1, vd2);
    }

    #[test]
    fn finished_verify_data_different_for_client_server() {
        let mut ks = KeySchedule::new(CipherSuite::Aes128Gcm, None);
        let shared_secret = vec![0x55u8; 32];
        let transcript = we_crypto::sha2::sha256(b"test");
        ks.derive_handshake_secrets(&shared_secret, &transcript);

        let client_vd = ks.compute_finished_verify_data(
            ks.client_handshake_traffic_secret().unwrap(),
            &transcript,
        );
        let server_vd = ks.compute_finished_verify_data(
            ks.server_handshake_traffic_secret().unwrap(),
            &transcript,
        );
        assert_ne!(client_vd, server_vd);
    }

    // -- RFC 8446 §A test vector (SHA-256 suite) --
    // Using constructed vectors to verify the key schedule structure is correct.

    #[test]
    fn key_schedule_full_flow_aes128() {
        let mut ks = KeySchedule::new(CipherSuite::Aes128Gcm, None);

        // Simulate handshake
        let mut transcript = TranscriptHash::new(CipherSuite::Aes128Gcm);
        transcript.update(b"ClientHello");
        transcript.update(b"ServerHello");
        let hs_hash = transcript.current_hash();

        let shared_secret = vec![0x01u8; 32]; // ECDHE shared secret
        ks.derive_handshake_secrets(&shared_secret, &hs_hash);

        // Get handshake keys
        let c_hs_keys = ks.client_handshake_keys();
        let s_hs_keys = ks.server_handshake_keys();
        assert_eq!(c_hs_keys.key.len(), 16);
        assert_eq!(s_hs_keys.key.len(), 16);

        // Continue transcript
        transcript.update(b"EncryptedExtensions");
        transcript.update(b"Certificate");
        transcript.update(b"CertificateVerify");

        // Server Finished
        let server_finished_hash = transcript.current_hash();
        let s_vd = ks.compute_finished_verify_data(
            ks.server_handshake_traffic_secret().unwrap(),
            &server_finished_hash,
        );
        assert_eq!(s_vd.len(), 32);

        transcript.update(b"Finished");
        let app_hash = transcript.current_hash();

        // Derive application secrets
        ks.derive_app_secrets(&app_hash);
        let c_app_keys = ks.client_app_keys();
        let s_app_keys = ks.server_app_keys();
        assert_eq!(c_app_keys.key.len(), 16);
        assert_eq!(s_app_keys.key.len(), 16);
        assert_ne!(c_app_keys.key, s_app_keys.key);
    }

    #[test]
    fn key_schedule_full_flow_aes256() {
        let mut ks = KeySchedule::new(CipherSuite::Aes256Gcm, None);
        let mut transcript = TranscriptHash::new(CipherSuite::Aes256Gcm);
        transcript.update(b"ClientHello");
        transcript.update(b"ServerHello");

        let shared_secret = vec![0x01u8; 48];
        ks.derive_handshake_secrets(&shared_secret, &transcript.current_hash());

        let keys = ks.client_handshake_keys();
        assert_eq!(keys.key.len(), 32);
        assert_eq!(keys.iv.len(), 12);

        transcript.update(b"rest of handshake");
        ks.derive_app_secrets(&transcript.current_hash());

        let app_keys = ks.client_app_keys();
        assert_eq!(app_keys.key.len(), 32);
    }

    #[test]
    fn key_schedule_full_flow_chacha20() {
        let mut ks = KeySchedule::new(CipherSuite::Chacha20Poly1305, None);
        let mut transcript = TranscriptHash::new(CipherSuite::Chacha20Poly1305);
        transcript.update(b"ClientHello");
        transcript.update(b"ServerHello");

        let shared_secret = vec![0x01u8; 32];
        ks.derive_handshake_secrets(&shared_secret, &transcript.current_hash());

        let keys = ks.server_handshake_keys();
        assert_eq!(keys.key.len(), 32);
        assert_eq!(keys.iv.len(), 12);

        transcript.update(b"rest of handshake");
        ks.derive_app_secrets(&transcript.current_hash());

        let app_keys = ks.server_app_keys();
        assert_eq!(app_keys.key.len(), 32);
    }

    // -- Derive-Secret tests --

    #[test]
    fn derive_secret_different_labels() {
        let secret = vec![0x42u8; 32];
        let ctx = we_crypto::sha2::sha256(b"context");
        let s1 = derive_secret::<Sha256>(&secret, b"label1", &ctx);
        let s2 = derive_secret::<Sha256>(&secret, b"label2", &ctx);
        assert_ne!(s1, s2);
    }

    #[test]
    fn derive_secret_sha384() {
        let secret = vec![0x42u8; 48];
        let ctx = we_crypto::sha2::sha384(b"context");
        let result = derive_secret::<Sha384>(&secret, b"test", &ctx);
        assert_eq!(result.len(), 48);
    }

    // -- Helper function tests --

    #[test]
    fn hash_len_for_suite_correct() {
        assert_eq!(hash_len_for_suite(CipherSuite::Aes128Gcm), 32);
        assert_eq!(hash_len_for_suite(CipherSuite::Aes256Gcm), 48);
        assert_eq!(hash_len_for_suite(CipherSuite::Chacha20Poly1305), 32);
    }

    #[test]
    fn empty_hash_correct_length() {
        assert_eq!(empty_hash(CipherSuite::Aes128Gcm).len(), 32);
        assert_eq!(empty_hash(CipherSuite::Aes256Gcm).len(), 48);
        assert_eq!(empty_hash(CipherSuite::Chacha20Poly1305).len(), 32);
    }

    #[test]
    fn empty_hash_matches_sha256_empty() {
        let result = empty_hash(CipherSuite::Aes128Gcm);
        let expected = we_crypto::sha2::sha256(&[]).to_vec();
        assert_eq!(result, expected);
    }

    #[test]
    fn empty_hash_matches_sha384_empty() {
        let result = empty_hash(CipherSuite::Aes256Gcm);
        let expected = we_crypto::sha2::sha384(&[]).to_vec();
        assert_eq!(result, expected);
    }
}