//! HKDF: HMAC-based Extract-and-Expand Key Derivation Function (RFC 5869).

use crate::hmac::{HashFunction, Hmac};

/// HKDF-Extract: derive a pseudorandom key (PRK) from input keying material.
///
/// `salt` is an optional non-secret random value; if empty, a string of
/// `H::OUTPUT_SIZE` zeros is used (per RFC 5869 §2.2).
pub fn hkdf_extract<H: HashFunction>(salt: &[u8], ikm: &[u8]) -> Vec<u8> {
    let effective_salt: Vec<u8>;
    let salt = if salt.is_empty() {
        effective_salt = vec![0u8; H::OUTPUT_SIZE];
        &effective_salt
    } else {
        salt
    };
    let mut hmac = Hmac::<H>::new(salt);
    hmac.update(ikm);
    hmac.finalize()
}

/// HKDF-Expand: expand a PRK into output keying material of length `len`.
///
/// `len` must be <= 255 * `H::OUTPUT_SIZE`.
/// Returns `None` if `len` exceeds the maximum.
pub fn hkdf_expand<H: HashFunction>(prk: &[u8], info: &[u8], len: usize) -> Option<Vec<u8>> {
    let hash_len = H::OUTPUT_SIZE;
    if len > 255 * hash_len {
        return None;
    }

    let n = len.div_ceil(hash_len);
    let mut okm = Vec::with_capacity(n * hash_len);
    let mut t_prev: Vec<u8> = Vec::new();

    for i in 1..=n {
        let mut hmac = Hmac::<H>::new(prk);
        hmac.update(&t_prev);
        hmac.update(info);
        hmac.update(&[i as u8]);
        t_prev = hmac.finalize();
        okm.extend_from_slice(&t_prev);
    }

    okm.truncate(len);
    Some(okm)
}

/// Combined HKDF: extract then expand in one call.
///
/// Returns `None` if `len` exceeds 255 * `H::OUTPUT_SIZE`.
pub fn hkdf<H: HashFunction>(salt: &[u8], ikm: &[u8], info: &[u8], len: usize) -> Option<Vec<u8>> {
    let prk = hkdf_extract::<H>(salt, ikm);
    hkdf_expand::<H>(&prk, info, len)
}

// ---------------------------------------------------------------------------
// Tests — RFC 5869 test vectors
// ---------------------------------------------------------------------------

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

    fn hex(bytes: &[u8]) -> String {
        bytes.iter().map(|b| format!("{b:02x}")).collect()
    }

    fn from_hex(s: &str) -> Vec<u8> {
        (0..s.len())
            .step_by(2)
            .map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
            .collect()
    }

    // -----------------------------------------------------------------------
    // Test Case 1: Basic test case with SHA-256
    // -----------------------------------------------------------------------

    #[test]
    fn rfc5869_case1_extract() {
        let ikm = from_hex("0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b");
        let salt = from_hex("000102030405060708090a0b0c");
        let prk = hkdf_extract::<Sha256>(&salt, &ikm);
        assert_eq!(
            hex(&prk),
            "077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e5"
        );
    }

    #[test]
    fn rfc5869_case1_expand() {
        let prk = from_hex("077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e5");
        let info = from_hex("f0f1f2f3f4f5f6f7f8f9");
        let okm = hkdf_expand::<Sha256>(&prk, &info, 42).unwrap();
        assert_eq!(
            hex(&okm),
            "3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865"
        );
    }

    #[test]
    fn rfc5869_case1_combined() {
        let ikm = from_hex("0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b");
        let salt = from_hex("000102030405060708090a0b0c");
        let info = from_hex("f0f1f2f3f4f5f6f7f8f9");
        let okm = hkdf::<Sha256>(&salt, &ikm, &info, 42).unwrap();
        assert_eq!(
            hex(&okm),
            "3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865"
        );
    }

    // -----------------------------------------------------------------------
    // Test Case 2: Longer inputs/outputs with SHA-256
    // -----------------------------------------------------------------------

    #[test]
    fn rfc5869_case2_extract() {
        let ikm = from_hex(
            "000102030405060708090a0b0c0d0e0f\
             101112131415161718191a1b1c1d1e1f\
             202122232425262728292a2b2c2d2e2f\
             303132333435363738393a3b3c3d3e3f\
             404142434445464748494a4b4c4d4e4f",
        );
        let salt = from_hex(
            "606162636465666768696a6b6c6d6e6f\
             707172737475767778797a7b7c7d7e7f\
             808182838485868788898a8b8c8d8e8f\
             909192939495969798999a9b9c9d9e9f\
             a0a1a2a3a4a5a6a7a8a9aaabacadaeaf",
        );
        let prk = hkdf_extract::<Sha256>(&salt, &ikm);
        assert_eq!(
            hex(&prk),
            "06a6b88c5853361a06104c9ceb35b45cef760014904671014a193f40c15fc244"
        );
    }

    #[test]
    fn rfc5869_case2_expand() {
        let prk = from_hex("06a6b88c5853361a06104c9ceb35b45cef760014904671014a193f40c15fc244");
        let info = from_hex(
            "b0b1b2b3b4b5b6b7b8b9babbbcbdbebf\
             c0c1c2c3c4c5c6c7c8c9cacbcccdcecf\
             d0d1d2d3d4d5d6d7d8d9dadbdcdddedf\
             e0e1e2e3e4e5e6e7e8e9eaebecedeeef\
             f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff",
        );
        let okm = hkdf_expand::<Sha256>(&prk, &info, 82).unwrap();
        assert_eq!(
            hex(&okm),
            "b11e398dc80327a1c8e7f78c596a4934\
             4f012eda2d4efad8a050cc4c19afa97c\
             59045a99cac7827271cb41c65e590e09\
             da3275600c2f09b8367793a9aca3db71\
             cc30c58179ec3e87c14c01d5c1f3434f\
             1d87"
        );
    }

    #[test]
    fn rfc5869_case2_combined() {
        let ikm = from_hex(
            "000102030405060708090a0b0c0d0e0f\
             101112131415161718191a1b1c1d1e1f\
             202122232425262728292a2b2c2d2e2f\
             303132333435363738393a3b3c3d3e3f\
             404142434445464748494a4b4c4d4e4f",
        );
        let salt = from_hex(
            "606162636465666768696a6b6c6d6e6f\
             707172737475767778797a7b7c7d7e7f\
             808182838485868788898a8b8c8d8e8f\
             909192939495969798999a9b9c9d9e9f\
             a0a1a2a3a4a5a6a7a8a9aaabacadaeaf",
        );
        let info = from_hex(
            "b0b1b2b3b4b5b6b7b8b9babbbcbdbebf\
             c0c1c2c3c4c5c6c7c8c9cacbcccdcecf\
             d0d1d2d3d4d5d6d7d8d9dadbdcdddedf\
             e0e1e2e3e4e5e6e7e8e9eaebecedeeef\
             f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff",
        );
        let okm = hkdf::<Sha256>(&salt, &ikm, &info, 82).unwrap();
        assert_eq!(
            hex(&okm),
            "b11e398dc80327a1c8e7f78c596a4934\
             4f012eda2d4efad8a050cc4c19afa97c\
             59045a99cac7827271cb41c65e590e09\
             da3275600c2f09b8367793a9aca3db71\
             cc30c58179ec3e87c14c01d5c1f3434f\
             1d87"
        );
    }

    // -----------------------------------------------------------------------
    // Test Case 3: SHA-256, zero-length salt and info
    // -----------------------------------------------------------------------

    #[test]
    fn rfc5869_case3_extract() {
        let ikm = from_hex("0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b");
        let prk = hkdf_extract::<Sha256>(&[], &ikm);
        assert_eq!(
            hex(&prk),
            "19ef24a32c717b167f33a91d6f648bdf96596776afdb6377ac434c1c293ccb04"
        );
    }

    #[test]
    fn rfc5869_case3_expand() {
        let prk = from_hex("19ef24a32c717b167f33a91d6f648bdf96596776afdb6377ac434c1c293ccb04");
        let okm = hkdf_expand::<Sha256>(&prk, &[], 42).unwrap();
        assert_eq!(
            hex(&okm),
            "8da4e775a563c18f715f802a063c5a31b8a11f5c5ee1879ec3454e5f3c738d2d9d201395faa4b61a96c8"
        );
    }

    #[test]
    fn rfc5869_case3_combined() {
        let ikm = from_hex("0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b");
        let okm = hkdf::<Sha256>(&[], &ikm, &[], 42).unwrap();
        assert_eq!(
            hex(&okm),
            "8da4e775a563c18f715f802a063c5a31b8a11f5c5ee1879ec3454e5f3c738d2d9d201395faa4b61a96c8"
        );
    }

    // -----------------------------------------------------------------------
    // Output length validation
    // -----------------------------------------------------------------------

    #[test]
    fn expand_rejects_oversized_output() {
        let prk = [0x42u8; 32];
        // 255 * 32 = 8160 is the max for SHA-256
        assert!(hkdf_expand::<Sha256>(&prk, &[], 8160).is_some());
        assert!(hkdf_expand::<Sha256>(&prk, &[], 8161).is_none());
    }

    #[test]
    fn hkdf_rejects_oversized_output() {
        let ikm = [0x0bu8; 22];
        assert!(hkdf::<Sha256>(&[], &ikm, &[], 8161).is_none());
    }

    // -----------------------------------------------------------------------
    // Edge cases
    // -----------------------------------------------------------------------

    #[test]
    fn expand_zero_length_output() {
        let prk = [0x42u8; 32];
        let okm = hkdf_expand::<Sha256>(&prk, &[], 0).unwrap();
        assert!(okm.is_empty());
    }

    #[test]
    fn expand_exact_hash_length() {
        let prk = [0x42u8; 32];
        let okm = hkdf_expand::<Sha256>(&prk, b"info", 32).unwrap();
        assert_eq!(okm.len(), 32);
    }

    #[test]
    fn extract_expand_with_sha512() {
        use crate::sha2::Sha512;

        // Use Test Case 1 inputs but with SHA-512 — verify it produces
        // the correct output length and doesn't panic.
        let ikm = from_hex("0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b");
        let salt = from_hex("000102030405060708090a0b0c");
        let prk = hkdf_extract::<Sha512>(&salt, &ikm);
        assert_eq!(prk.len(), 64);

        let okm = hkdf_expand::<Sha512>(&prk, b"test info", 100).unwrap();
        assert_eq!(okm.len(), 100);
    }

    #[test]
    fn extract_expand_with_sha384() {
        use crate::sha2::Sha384;

        let ikm = [0xaau8; 80];
        let prk = hkdf_extract::<Sha384>(&[], &ikm);
        assert_eq!(prk.len(), 48);

        let okm = hkdf_expand::<Sha384>(&prk, b"context", 60).unwrap();
        assert_eq!(okm.len(), 60);
    }
}