//! HKDF (HMAC-based Key Derivation Function) implementation
//!
//! Implements RFC 5869 HKDF with SHA-256 and SHA-512 hash functions.
//! Used for deriving session keys and sync tokens.

use hkdf::Hkdf;
use sha2::{Sha256, Sha512};
use zeroize::Zeroizing;

use crate::error::{CryptoError, CryptoResult};

/// Maximum output length for HKDF-SHA256 (255 * 32 = 8160 bytes)
pub const MAX_OUTPUT_SHA256: usize = 255 * 32;

/// Maximum output length for HKDF-SHA512 (255 * 64 = 16320 bytes)
pub const MAX_OUTPUT_SHA512: usize = 255 * 64;

/// Derive a key using HKDF-SHA256
///
/// # Arguments
/// * `secret` - Input key material (IKM)
/// * `salt` - Optional salt value (recommended for security)
/// * `info` - Context and application-specific information
/// * `length` - Desired output key length in bytes (max 8160)
///
/// # Returns
/// Derived key bytes, or error if length exceeds maximum
///
/// # Example
/// ```
/// use pqc_py::keys::hkdf::derive_key;
///
/// let secret = b"my secret key material";
/// let salt = Some(b"optional salt".as_slice());
/// let info = b"pqc-session-v1";
///
/// let derived = derive_key(secret, salt, info, 32).unwrap();
/// assert_eq!(derived.len(), 32);
/// ```
pub fn derive_key(
    secret: &[u8],
    salt: Option<&[u8]>,
    info: &[u8],
    length: usize,
) -> CryptoResult<Vec<u8>> {
    if length > MAX_OUTPUT_SHA256 {
        return Err(CryptoError::InvalidOutputLength(format!(
            "Requested {} bytes, maximum is {}",
            length, MAX_OUTPUT_SHA256
        )));
    }

    if length == 0 {
        return Err(CryptoError::InvalidOutputLength(
            "Output length must be greater than 0".to_string()
        ));
    }

    let hk = Hkdf::<Sha256>::new(salt, secret);
    let mut okm = Zeroizing::new(vec![0u8; length]);

    hk.expand(info, &mut okm)
        .map_err(|e| CryptoError::KeyDerivationFailed(e.to_string()))?;

    Ok(okm.to_vec())
}

/// Derive a key using HKDF-SHA512
///
/// Provides larger output capacity than SHA256 variant.
///
/// # Arguments
/// * `secret` - Input key material (IKM)
/// * `salt` - Optional salt value (recommended for security)
/// * `info` - Context and application-specific information
/// * `length` - Desired output key length in bytes (max 16320)
///
/// # Returns
/// Derived key bytes, or error if length exceeds maximum
pub fn derive_key_sha512(
    secret: &[u8],
    salt: Option<&[u8]>,
    info: &[u8],
    length: usize,
) -> CryptoResult<Vec<u8>> {
    if length > MAX_OUTPUT_SHA512 {
        return Err(CryptoError::InvalidOutputLength(format!(
            "Requested {} bytes, maximum is {}",
            length, MAX_OUTPUT_SHA512
        )));
    }

    if length == 0 {
        return Err(CryptoError::InvalidOutputLength(
            "Output length must be greater than 0".to_string()
        ));
    }

    let hk = Hkdf::<Sha512>::new(salt, secret);
    let mut okm = Zeroizing::new(vec![0u8; length]);

    hk.expand(info, &mut okm)
        .map_err(|e| CryptoError::KeyDerivationFailed(e.to_string()))?;

    Ok(okm.to_vec())
}

/// Derive a fixed 32-byte key using HKDF-SHA256
///
/// Convenience function for the common case of deriving a 256-bit key.
pub fn derive_key_32(
    secret: &[u8],
    salt: Option<&[u8]>,
    info: &[u8],
) -> CryptoResult<[u8; 32]> {
    let derived = derive_key(secret, salt, info, 32)?;
    let mut key = [0u8; 32];
    key.copy_from_slice(&derived);
    Ok(key)
}

/// Derive a fixed 64-byte key using HKDF-SHA512
///
/// Convenience function for deriving a 512-bit key.
pub fn derive_key_64(
    secret: &[u8],
    salt: Option<&[u8]>,
    info: &[u8],
) -> CryptoResult<[u8; 64]> {
    let derived = derive_key_sha512(secret, salt, info, 64)?;
    let mut key = [0u8; 64];
    key.copy_from_slice(&derived);
    Ok(key)
}

/// Derive a session key for a covert channel
///
/// Uses a standardized info format for session key derivation.
///
/// # Arguments
/// * `master_key` - The master key material
/// * `session_id` - Unique session identifier
/// * `channel` - Channel name (e.g., "TCOE", "CoTSE")
pub fn derive_session_key(
    master_key: &[u8],
    session_id: &str,
    channel: &str,
) -> CryptoResult<[u8; 32]> {
    let info = format!("phantom-protocol:session:{}:{}", channel, session_id);
    derive_key_32(master_key, None, info.as_bytes())
}

/// Derive a sync token for cross-agent synchronization
///
/// # Arguments
/// * `shared_key` - The shared key between agents
/// * `timestamp` - Unix timestamp for token validity
/// * `agent_ids` - Sorted list of participating agent IDs
pub fn derive_sync_token(
    shared_key: &[u8],
    timestamp: u64,
    agent_ids: &[&str],
) -> CryptoResult<[u8; 16]> {
    let agents = agent_ids.join(",");
    let info = format!("phantom-protocol:sync:{}:{}", timestamp, agents);
    let derived = derive_key(shared_key, None, info.as_bytes(), 16)?;
    let mut token = [0u8; 16];
    token.copy_from_slice(&derived);
    Ok(token)
}

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

    #[test]
    fn test_derive_key_basic() {
        let secret = b"test secret";
        let salt = Some(b"test salt".as_slice());
        let info = b"test info";

        let key1 = derive_key(secret, salt, info, 32).unwrap();
        let key2 = derive_key(secret, salt, info, 32).unwrap();

        assert_eq!(key1.len(), 32);
        assert_eq!(key1, key2); // Same inputs produce same output
    }

    #[test]
    fn test_derive_key_different_salt() {
        let secret = b"test secret";
        let info = b"test info";

        let key1 = derive_key(secret, Some(b"salt1"), info, 32).unwrap();
        let key2 = derive_key(secret, Some(b"salt2"), info, 32).unwrap();

        assert_ne!(key1, key2); // Different salt produces different keys
    }

    #[test]
    fn test_derive_key_different_info() {
        let secret = b"test secret";
        let salt = Some(b"salt".as_slice());

        let key1 = derive_key(secret, salt, b"info1", 32).unwrap();
        let key2 = derive_key(secret, salt, b"info2", 32).unwrap();

        assert_ne!(key1, key2); // Different info produces different keys
    }

    #[test]
    fn test_derive_key_no_salt() {
        let secret = b"test secret";
        let info = b"test info";

        let key = derive_key(secret, None, info, 32).unwrap();
        assert_eq!(key.len(), 32);
    }

    #[test]
    fn test_derive_key_various_lengths() {
        let secret = b"test secret";
        let salt = Some(b"salt".as_slice());
        let info = b"info";

        for length in [16, 32, 48, 64, 128, 256] {
            let key = derive_key(secret, salt, info, length).unwrap();
            assert_eq!(key.len(), length);
        }
    }

    #[test]
    fn test_derive_key_max_length() {
        let secret = b"test secret";
        let key = derive_key(secret, None, b"info", MAX_OUTPUT_SHA256).unwrap();
        assert_eq!(key.len(), MAX_OUTPUT_SHA256);
    }

    #[test]
    fn test_derive_key_exceeds_max() {
        let secret = b"test secret";
        let result = derive_key(secret, None, b"info", MAX_OUTPUT_SHA256 + 1);
        assert!(matches!(result, Err(CryptoError::InvalidOutputLength(_))));
    }

    #[test]
    fn test_derive_key_zero_length() {
        let secret = b"test secret";
        let result = derive_key(secret, None, b"info", 0);
        assert!(matches!(result, Err(CryptoError::InvalidOutputLength(_))));
    }

    #[test]
    fn test_derive_key_sha512() {
        let secret = b"test secret";
        let salt = Some(b"salt".as_slice());
        let info = b"info";

        let key = derive_key_sha512(secret, salt, info, 64).unwrap();
        assert_eq!(key.len(), 64);
    }

    #[test]
    fn test_derive_key_32() {
        let secret = b"test secret";
        let key = derive_key_32(secret, None, b"info").unwrap();
        assert_eq!(key.len(), 32);
    }

    #[test]
    fn test_derive_key_64() {
        let secret = b"test secret";
        let key = derive_key_64(secret, None, b"info").unwrap();
        assert_eq!(key.len(), 64);
    }

    #[test]
    fn test_derive_session_key() {
        let master = b"master key material";
        let key = derive_session_key(master, "session-123", "TCOE").unwrap();
        assert_eq!(key.len(), 32);
    }

    #[test]
    fn test_derive_sync_token() {
        let shared = b"shared key";
        let token = derive_sync_token(shared, 1703980800, &["agent-1", "agent-2"]).unwrap();
        assert_eq!(token.len(), 16);
    }

    #[test]
    fn test_derive_sync_token_deterministic() {
        let shared = b"shared key";
        let agents = ["agent-1", "agent-2"];
        let timestamp = 1703980800u64;

        let token1 = derive_sync_token(shared, timestamp, &agents).unwrap();
        let token2 = derive_sync_token(shared, timestamp, &agents).unwrap();

        assert_eq!(token1, token2);
    }

    // Test vector from RFC 5869 Appendix A.1
    #[test]
    fn test_rfc5869_test_vector() {
        let ikm = hex::decode("0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b").unwrap();
        let salt = hex::decode("000102030405060708090a0b0c").unwrap();
        let info = hex::decode("f0f1f2f3f4f5f6f7f8f9").unwrap();

        let expected = hex::decode(
            "3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865"
        ).unwrap();

        let okm = derive_key(&ikm, Some(&salt), &info, 42).unwrap();
        assert_eq!(okm, expected);
    }
}
