//! Vault-controlled key derivation for NTAG 424 DNA tag personalization.
//!
//! Derives per-tag AES-128 key sets deterministically from a vault root secret
//! using HKDF-SHA256. Supports key versioning for controlled rotation.
//! Per the key management approach in docs/nfc.md §10.

use hkdf::Hkdf;
use sha2::Sha256;
use zeroize::Zeroize;

/// The five AES-128 keys for an NTAG 424 DNA application.
/// Index matches key number: 0=AppMaster, 1=SDMMeta, 2=SDMFile, 3=Config, 4=Reserve.
#[derive(Clone)]
pub struct TagKeySet {
    pub keys: [[u8; 16]; 5],
    pub key_version: u32,
    pub tag_uid: Vec<u8>,
}

impl std::fmt::Debug for TagKeySet {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TagKeySet")
            .field("key_version", &self.key_version)
            .field("tag_uid", &hex::encode(&self.tag_uid))
            .field("keys", &"[REDACTED]")
            .finish()
    }
}

impl Drop for TagKeySet {
    fn drop(&mut self) {
        for key in &mut self.keys {
            key.zeroize();
        }
    }
}

/// Key purpose labels matching the house convention in the implementation spec.
const KEY_LABELS: [&str; 5] = [
    "ntag424/master",
    "ntag424/sdm_meta",
    "ntag424/sdm_file",
    "ntag424/config",
    "ntag424/reserve",
];

/// Derive a complete set of 5 AES-128 keys for a specific tag.
///
/// `vault_secret`: root secret from the vault (any length, should be high-entropy).
/// `tag_uid`: 7-byte UID from the tag.
/// `key_version`: monotonic version for key rotation; incrementing this produces
///                a completely new key set.
///
/// Derivation: for each key index i (0..5):
///   salt = tag_uid || key_version (LE 4 bytes)
///   K_i  = HKDF-SHA256(ikm=vault_secret, salt=salt, info=KEY_LABELS[i])
pub fn derive_tag_keys(
    vault_secret: &[u8],
    tag_uid: &[u8],
    key_version: u32,
) -> TagKeySet {
    let mut salt = Vec::with_capacity(tag_uid.len() + 4);
    salt.extend_from_slice(tag_uid);
    salt.extend_from_slice(&key_version.to_le_bytes());

    let hkdf = Hkdf::<Sha256>::new(Some(&salt), vault_secret);

    let mut keys = [[0u8; 16]; 5];
    for (i, label) in KEY_LABELS.iter().enumerate() {
        hkdf.expand(label.as_bytes(), &mut keys[i])
            .expect("HKDF expand should not fail for 16-byte output");
    }

    TagKeySet {
        keys,
        key_version,
        tag_uid: tag_uid.to_vec(),
    }
}

/// Derive a single purpose key (for cases where only one key is needed).
pub fn derive_single_key(
    vault_secret: &[u8],
    tag_uid: &[u8],
    key_version: u32,
    key_index: usize,
) -> Result<[u8; 16], String> {
    if key_index >= 5 {
        return Err(format!("key_index must be 0..4, got {key_index}"));
    }
    let mut salt = Vec::with_capacity(tag_uid.len() + 4);
    salt.extend_from_slice(tag_uid);
    salt.extend_from_slice(&key_version.to_le_bytes());

    let hkdf = Hkdf::<Sha256>::new(Some(&salt), vault_secret);
    let mut key = [0u8; 16];
    hkdf.expand(KEY_LABELS[key_index].as_bytes(), &mut key)
        .map_err(|e| format!("HKDF expand failed: {e}"))?;
    Ok(key)
}

/// The all-zero transport key used on factory-default tags.
pub const TRANSPORT_KEY: [u8; 16] = [0x00; 16];

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

    #[test]
    fn derive_deterministic() {
        let secret = b"test-vault-secret-32-bytes-long!";
        let uid = [0x04, 0xAA, 0xBB, 0xCC, 0xDD, 0x11, 0x22];
        let ks1 = derive_tag_keys(secret, &uid, 1);
        let ks2 = derive_tag_keys(secret, &uid, 1);
        assert_eq!(ks1.keys, ks2.keys);
    }

    #[test]
    fn different_versions_produce_different_keys() {
        let secret = b"test-vault-secret";
        let uid = [0x04, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06];
        let ks1 = derive_tag_keys(secret, &uid, 1);
        let ks2 = derive_tag_keys(secret, &uid, 2);
        assert_ne!(ks1.keys[0], ks2.keys[0]);
    }

    #[test]
    fn different_uids_produce_different_keys() {
        let secret = b"test-vault-secret";
        let uid_a = [0x04, 0xAA, 0xBB, 0xCC, 0xDD, 0x11, 0x22];
        let uid_b = [0x04, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66];
        let ks_a = derive_tag_keys(secret, &uid_a, 1);
        let ks_b = derive_tag_keys(secret, &uid_b, 1);
        assert_ne!(ks_a.keys[0], ks_b.keys[0]);
    }

    #[test]
    fn all_five_keys_are_unique() {
        let secret = b"another-secret";
        let uid = [0x04, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06];
        let ks = derive_tag_keys(secret, &uid, 0);
        for i in 0..5 {
            for j in (i + 1)..5 {
                assert_ne!(ks.keys[i], ks.keys[j], "key {i} == key {j}");
            }
        }
    }

    #[test]
    fn single_key_matches_full_set() {
        let secret = b"consistency-test";
        let uid = [0x04, 0xDE, 0xAD, 0xBE, 0xEF, 0x01, 0x02];
        let ks = derive_tag_keys(secret, &uid, 3);
        for i in 0..5 {
            let single = derive_single_key(secret, &uid, 3, i).unwrap();
            assert_eq!(single, ks.keys[i]);
        }
    }

}
