//! NTAG 424 DNA AuthenticateEV2First session management.
//!
//! Implements the two-part EV2 authentication exchange, session key derivation
//! (KSesAuthENC / KSesAuthMAC), command MAC computation, and CommMode.Full
//! encryption/decryption. Per NXP AN12196.

use aes::cipher::{BlockEncrypt, KeyInit};
use aes::Aes128;
use cmac::{Cmac, Mac};

/// Active EV2 session state. Persists across commands until the session is
/// invalidated (tag removal, error, or explicit re-auth).
#[derive(Debug, Clone)]
pub struct Ev2Session {
    pub auth_key_no: u8,
    /// Transaction Identifier (4 bytes), assigned by PICC.
    pub ti: [u8; 4],
    /// Command counter, incremented before each authenticated command.
    pub cmd_ctr: u16,
    /// Session encryption key.
    pub ks_enc: [u8; 16],
    /// Session MAC key.
    pub ks_mac: [u8; 16],
}

// ---------------------------------------------------------------------------
// AES primitives (block-aligned data only — NTAG 424 DNA invariant)
// ---------------------------------------------------------------------------

fn aes_ecb_encrypt_block(key: &[u8; 16], block: &[u8; 16]) -> [u8; 16] {
    let cipher = Aes128::new(key.into());
    let mut out = *block;
    cipher.encrypt_block((&mut out).into());
    out
}

fn aes_cbc_encrypt(key: &[u8; 16], iv: &[u8; 16], plaintext: &[u8]) -> Vec<u8> {
    assert!(
        plaintext.len() % 16 == 0 && !plaintext.is_empty(),
        "AES-CBC plaintext must be non-empty and block-aligned"
    );
    let cipher = Aes128::new(key.into());
    let mut result = Vec::with_capacity(plaintext.len());
    let mut prev = *iv;
    for chunk in plaintext.chunks_exact(16) {
        let mut block = [0u8; 16];
        for i in 0..16 {
            block[i] = chunk[i] ^ prev[i];
        }
        cipher.encrypt_block((&mut block).into());
        prev = block;
        result.extend_from_slice(&block);
    }
    result
}

fn aes_cbc_decrypt(key: &[u8; 16], iv: &[u8; 16], ciphertext: &[u8]) -> Result<Vec<u8>, String> {
    if ciphertext.len() % 16 != 0 || ciphertext.is_empty() {
        return Err("AES-CBC ciphertext must be non-empty and block-aligned".into());
    }
    use aes::cipher::BlockDecrypt;
    let decipher = Aes128::new(key.into());
    let mut result = Vec::with_capacity(ciphertext.len());
    let mut prev_block = *iv;
    for chunk in ciphertext.chunks_exact(16) {
        let mut block = [0u8; 16];
        block.copy_from_slice(chunk);
        decipher.decrypt_block((&mut block).into());
        for i in 0..16 {
            block[i] ^= prev_block[i];
        }
        prev_block.copy_from_slice(chunk);
        result.extend_from_slice(&block);
    }
    Ok(result)
}

/// Compute full AES-128-CMAC (16 bytes).
fn aes_cmac(key: &[u8; 16], data: &[u8]) -> [u8; 16] {
    let mut mac = <Cmac<Aes128> as Mac>::new(key.into());
    mac.update(data);
    let result = mac.finalize().into_bytes();
    let mut out = [0u8; 16];
    out.copy_from_slice(&result);
    out
}

/// Truncate a 16-byte CMAC to 8 bytes for MACt.
/// Per Enjoy RFID / AN12196: odd-indexed bytes (1,3,5,7,9,11,13,15).
fn truncate_mac(full_mac: &[u8; 16]) -> [u8; 8] {
    let mut t = [0u8; 8];
    for i in 0..8 {
        t[i] = full_mac[2 * i + 1];
    }
    t
}

/// Rotate a 16-byte array left by 1 byte: [a,b,c,...,p] → [b,c,...,p,a].
fn rotate_left_1(data: &[u8; 16]) -> [u8; 16] {
    let mut out = [0u8; 16];
    out[..15].copy_from_slice(&data[1..]);
    out[15] = data[0];
    out
}

// ---------------------------------------------------------------------------
// Session key derivation (AN12196 §2.3.1)
// ---------------------------------------------------------------------------

/// Build Session Vector for key derivation.
/// `label`: 0xA55A for ENC, 0x5AA5 for MAC.
fn build_sv(label: [u8; 2], rnd_a: &[u8; 16], rnd_b: &[u8; 16]) -> [u8; 32] {
    let mut sv = [0u8; 32];
    sv[0] = label[0];
    sv[1] = label[1];
    sv[2] = 0x00;
    sv[3] = 0x01;
    sv[4] = 0x00;
    sv[5] = 0x80;
    // RndA[0..2]
    sv[6] = rnd_a[0];
    sv[7] = rnd_a[1];
    // RndA[2..8] XOR RndB[0..6]
    for i in 0..6 {
        sv[8 + i] = rnd_a[2 + i] ^ rnd_b[i];
    }
    // RndB[6..16]
    sv[14..24].copy_from_slice(&rnd_b[6..16]);
    // RndA[8..16]
    sv[24..32].copy_from_slice(&rnd_a[8..16]);
    sv
}

fn derive_session_keys(
    auth_key: &[u8; 16],
    rnd_a: &[u8; 16],
    rnd_b: &[u8; 16],
) -> ([u8; 16], [u8; 16]) {
    let sv1 = build_sv([0xA5, 0x5A], rnd_a, rnd_b);
    let sv2 = build_sv([0x5A, 0xA5], rnd_a, rnd_b);
    let ks_enc = aes_cmac(auth_key, &sv1);
    let ks_mac = aes_cmac(auth_key, &sv2);
    (ks_enc, ks_mac)
}

// ---------------------------------------------------------------------------
// AuthenticateEV2First protocol
// ---------------------------------------------------------------------------

/// Intermediate state after receiving the PICC's first response (encrypted RndB).
/// The caller must send part2 APDU and then call `complete_auth` with the response.
pub struct Ev2AuthChallenge {
    auth_key: [u8; 16],
    auth_key_no: u8,
    rnd_a: [u8; 16],
    rnd_b: [u8; 16],
    /// The APDU payload for part 2 (encrypted RndA || RndB').
    pub part2_payload: Vec<u8>,
}

/// Process part 1 response (encrypted RndB) and prepare part 2 payload.
/// Returns the challenge state needed to complete authentication.
///
/// `part1_data`: the payload from the 91AF response (16 bytes = encrypted RndB).
/// `auth_key`: the 16-byte AES key for the target key number.
/// `key_no`: the key number being authenticated.
pub fn process_auth_part1(
    part1_data: &[u8],
    auth_key: &[u8; 16],
    key_no: u8,
) -> Result<Ev2AuthChallenge, String> {
    if part1_data.len() != 16 {
        return Err(format!(
            "auth part1 response should be 16 bytes, got {}",
            part1_data.len()
        ));
    }

    let iv_zero = [0u8; 16];
    let rnd_b_vec = aes_cbc_decrypt(auth_key, &iv_zero, part1_data)?;
    let mut rnd_b = [0u8; 16];
    rnd_b.copy_from_slice(&rnd_b_vec);

    let mut rnd_a = [0u8; 16];
    use rand::RngCore;
    rand::rngs::OsRng.fill_bytes(&mut rnd_a);

    let rnd_b_rot = rotate_left_1(&rnd_b);
    let mut plaintext = [0u8; 32];
    plaintext[..16].copy_from_slice(&rnd_a);
    plaintext[16..].copy_from_slice(&rnd_b_rot);

    let ciphertext = aes_cbc_encrypt(auth_key, &iv_zero, &plaintext);

    Ok(Ev2AuthChallenge {
        auth_key: *auth_key,
        auth_key_no: key_no,
        rnd_a,
        rnd_b,
        part2_payload: ciphertext,
    })
}

/// Complete authentication using the PICC's part 2 response.
/// `part2_data`: payload from the 9100 response (32 bytes = encrypted TI || RndA' || PDcap2 || PCDcap2).
pub fn complete_auth(
    challenge: Ev2AuthChallenge,
    part2_data: &[u8],
) -> Result<Ev2Session, String> {
    if part2_data.len() != 32 {
        return Err(format!(
            "auth part2 response should be 32 bytes, got {}",
            part2_data.len()
        ));
    }

    let iv_zero = [0u8; 16];
    let decrypted = aes_cbc_decrypt(&challenge.auth_key, &iv_zero, part2_data)?;

    let mut ti = [0u8; 4];
    ti.copy_from_slice(&decrypted[0..4]);

    let mut rnd_a_prime = [0u8; 16];
    rnd_a_prime.copy_from_slice(&decrypted[4..20]);

    let expected_rnd_a_prime = rotate_left_1(&challenge.rnd_a);
    if rnd_a_prime != expected_rnd_a_prime {
        return Err("authentication failed: RndA' mismatch".into());
    }

    let (ks_enc, ks_mac) =
        derive_session_keys(&challenge.auth_key, &challenge.rnd_a, &challenge.rnd_b);

    Ok(Ev2Session {
        auth_key_no: challenge.auth_key_no,
        ti,
        cmd_ctr: 0,
        ks_enc,
        ks_mac,
    })
}

// ---------------------------------------------------------------------------
// Authenticated command helpers
// ---------------------------------------------------------------------------

impl Ev2Session {
    /// Compute the truncated 8-byte command MAC.
    /// `cmd`: the NXP command byte (e.g. 0xF5 for GetFileSettings).
    /// `cmd_header`: any command-specific header bytes after the command code.
    /// `cmd_data`: the (possibly encrypted) command data payload.
    pub fn command_mac(&self, cmd: u8, cmd_header: &[u8], cmd_data: &[u8]) -> [u8; 8] {
        let mut mac_input = Vec::new();
        mac_input.push(cmd);
        mac_input.extend_from_slice(&self.cmd_ctr.to_le_bytes());
        mac_input.extend_from_slice(&self.ti);
        mac_input.extend_from_slice(cmd_header);
        mac_input.extend_from_slice(cmd_data);
        let full = aes_cmac(&self.ks_mac, &mac_input);
        truncate_mac(&full)
    }

    /// Compute the truncated 8-byte response MAC for verification.
    /// `response_code`: the SW2 byte (e.g. 0x00 for success).
    /// `response_data`: the (possibly encrypted) response data before the MAC.
    pub fn response_mac(&self, response_code: u8, response_data: &[u8]) -> [u8; 8] {
        let mut mac_input = Vec::new();
        mac_input.push(response_code);
        mac_input.extend_from_slice(&self.cmd_ctr.to_le_bytes());
        mac_input.extend_from_slice(&self.ti);
        mac_input.extend_from_slice(response_data);
        let full = aes_cmac(&self.ks_mac, &mac_input);
        truncate_mac(&full)
    }

    /// Compute the IV for command encryption.
    fn command_iv(&self) -> [u8; 16] {
        let mut iv_input = [0u8; 16];
        iv_input[0] = 0xA5;
        iv_input[1] = 0x5A;
        iv_input[2..6].copy_from_slice(&self.ti);
        iv_input[6..8].copy_from_slice(&self.cmd_ctr.to_le_bytes());
        aes_ecb_encrypt_block(&self.ks_enc, &iv_input)
    }

    /// Compute the IV for response decryption.
    fn response_iv(&self) -> [u8; 16] {
        let mut iv_input = [0u8; 16];
        iv_input[0] = 0x5A;
        iv_input[1] = 0xA5;
        iv_input[2..6].copy_from_slice(&self.ti);
        iv_input[6..8].copy_from_slice(&self.cmd_ctr.to_le_bytes());
        aes_ecb_encrypt_block(&self.ks_enc, &iv_input)
    }

    /// Encrypt command data using CommMode.Full.
    /// Data is padded with 0x80 || 0x00... to the next 16-byte boundary.
    pub fn encrypt_command(&self, plaintext: &[u8]) -> Vec<u8> {
        let padded = pad_iso9797_m2(plaintext);
        let iv = self.command_iv();
        aes_cbc_encrypt(&self.ks_enc, &iv, &padded)
    }

    /// Decrypt response data from CommMode.Full.
    /// Removes ISO 9797-1 M2 padding after decryption.
    pub fn decrypt_response(&self, ciphertext: &[u8]) -> Result<Vec<u8>, String> {
        let iv = self.response_iv();
        let padded = aes_cbc_decrypt(&self.ks_enc, &iv, ciphertext)?;
        unpad_iso9797_m2(&padded)
    }

    /// Increment the command counter. Call before each authenticated command.
    pub fn increment_cmd_ctr(&mut self) {
        self.cmd_ctr = self.cmd_ctr.wrapping_add(1);
    }

    /// Build a MACed command (CommMode.MAC).
    /// Returns the full wrapped APDU ready to transmit.
    /// If `first_command_use_ctr_zero` is true, use CmdCtr=0 (required for first command after auth).
    pub fn build_maced_command(
        &mut self,
        cmd: u8,
        cmd_header: &[u8],
        cmd_data: &[u8],
        first_command_use_ctr_zero: bool,
    ) -> Vec<u8> {
        if first_command_use_ctr_zero {
            self.cmd_ctr = 0;
        } else {
            self.increment_cmd_ctr();
        }
        let mac = self.command_mac(cmd, cmd_header, cmd_data);
        if first_command_use_ctr_zero {
            self.cmd_ctr = 1;
        }
        let mut payload = Vec::new();
        payload.extend_from_slice(cmd_header);
        payload.extend_from_slice(cmd_data);
        payload.extend_from_slice(&mac);
        // Build wrapped APDU
        let mut apdu = Vec::with_capacity(6 + payload.len());
        apdu.push(0x90);
        apdu.push(cmd);
        apdu.push(0x00);
        apdu.push(0x00);
        apdu.push(payload.len() as u8);
        apdu.extend_from_slice(&payload);
        apdu.push(0x00);
        apdu
    }

    /// Build ChangeKey (C4h) in CommMode.Full. Per AN12196 Table 26: compact format
    /// (CmdHeader || Encrypted || MAC = 41 bytes). CmdCtr/TI not transmitted.
    /// first_ctr_zero: use CmdCtr=0 for first ChangeKey (chip may not increment for compact CFS).
    pub fn build_change_key_command(
        &mut self,
        cmd_header: &[u8],
        cmd_data: &[u8],
        first_ctr_zero: bool,
    ) -> Vec<u8> {
        if first_ctr_zero {
            self.cmd_ctr = 0;
        }
        let encrypted = self.encrypt_command(cmd_data);
        let mac = self.command_mac(0xC4, cmd_header, &encrypted);
        if first_ctr_zero {
            self.cmd_ctr = 1;
        } else {
            self.increment_cmd_ctr();
        }
        let mut payload = Vec::new();
        payload.extend_from_slice(cmd_header);
        payload.extend_from_slice(&encrypted);
        payload.extend_from_slice(&mac);
        let mut apdu = Vec::with_capacity(6 + payload.len());
        apdu.push(0x90);
        apdu.push(0xC4);
        apdu.push(0x00);
        apdu.push(0x00);
        apdu.push(payload.len() as u8);
        apdu.extend_from_slice(&payload);
        apdu.push(0x00);
        apdu
    }

    /// Build ChangeFileSettings (5Fh) in CommMode.Full.
    /// Per AN12196 §6.8 / Enjoy RFID: payload = CmdHeader || EncryptedCmdData || MACt (25 bytes).
    /// CmdCtr and TI are NOT transmitted; MAC is still computed over Cmd || CmdCtr || TI || CmdHeader || Encrypted.
    ///
    /// If `first_command_use_ctr_zero` is true, use CmdCtr=0 for IV/MAC (some readers expect 0 for first command).
    pub fn build_change_file_settings_command(
        &mut self,
        cmd_header: &[u8],
        cmd_data: &[u8],
        first_command_use_ctr_zero: bool,
    ) -> Vec<u8> {
        if first_command_use_ctr_zero {
            self.cmd_ctr = 0; // use 0 for first command
        } else {
            self.increment_cmd_ctr(); // 0 -> 1
        }
        let encrypted = self.encrypt_command(cmd_data);
        let mac = self.command_mac(0x5F, cmd_header, &encrypted);
        if first_command_use_ctr_zero {
            self.cmd_ctr = 1; // advance for next command
        }
        let mut payload = Vec::new();
        payload.extend_from_slice(cmd_header);
        payload.extend_from_slice(&encrypted);
        payload.extend_from_slice(&mac);
        let mut apdu = Vec::with_capacity(6 + payload.len());
        apdu.push(0x90);
        apdu.push(0x5F);
        apdu.push(0x00);
        apdu.push(0x00);
        apdu.push(payload.len() as u8);
        apdu.extend_from_slice(&payload);
        apdu.push(0x00);
        apdu
    }

    /// Verify response MAC, optionally overriding the cmd_ctr used (for first command after auth).
    pub fn verify_response_mac_with_ctr(
        &self,
        status_code: u8,
        response_data: &[u8],
        cmd_ctr_override: Option<u16>,
    ) -> Result<Vec<u8>, String> {
        if response_data.len() < 8 {
            return Err("response too short for MAC verification".into());
        }
        let data_len = response_data.len() - 8;
        let data = &response_data[..data_len];
        let received_mac = &response_data[data_len..];
        let expected = if let Some(ctr) = cmd_ctr_override {
            let mut mac_input = Vec::new();
            mac_input.push(status_code);
            mac_input.extend_from_slice(&ctr.to_le_bytes());
            mac_input.extend_from_slice(&self.ti);
            mac_input.extend_from_slice(data);
            let full = aes_cmac(&self.ks_mac, &mac_input);
            truncate_mac(&full)
        } else {
            self.response_mac(status_code, data)
        };
        if received_mac != expected {
            return Err("response MAC mismatch".into());
        }
        Ok(data.to_vec())
    }
}

// ---------------------------------------------------------------------------
// Padding (ISO 9797-1 Method 2)
// ---------------------------------------------------------------------------

/// Pad data per ISO 9797-1 Method 2: append 0x80 then zeros to next block boundary.
fn pad_iso9797_m2(data: &[u8]) -> Vec<u8> {
    let padded_len = ((data.len() + 1 + 15) / 16) * 16;
    let mut out = Vec::with_capacity(padded_len);
    out.extend_from_slice(data);
    out.push(0x80);
    out.resize(padded_len, 0x00);
    out
}

/// Remove ISO 9797-1 M2 padding.
fn unpad_iso9797_m2(data: &[u8]) -> Result<Vec<u8>, String> {
    for i in (0..data.len()).rev() {
        match data[i] {
            0x80 => return Ok(data[..i].to_vec()),
            0x00 => continue,
            _ => return Err("invalid ISO 9797-1 M2 padding".into()),
        }
    }
    Err("no padding marker found".into())
}

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

    #[test]
    fn rotate_left_1_basic() {
        let input: [u8; 16] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
        let rotated = rotate_left_1(&input);
        assert_eq!(
            rotated,
            [2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 1]
        );
    }

    #[test]
    fn aes_cbc_round_trip() {
        let key = [0x00u8; 16];
        let iv = [0x00u8; 16];
        let plaintext = [0x42u8; 32];
        let ct = aes_cbc_encrypt(&key, &iv, &plaintext);
        let pt = aes_cbc_decrypt(&key, &iv, &ct).unwrap();
        assert_eq!(pt, plaintext);
    }

    #[test]
    fn cmac_truncation() {
        let full: [u8; 16] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15];
        let t = truncate_mac(&full);
        assert_eq!(t, [1, 3, 5, 7, 9, 11, 13, 15]);
    }

    #[test]
    fn sv_construction_length() {
        let rnd_a = [0xAAu8; 16];
        let rnd_b = [0xBBu8; 16];
        let sv = build_sv([0xA5, 0x5A], &rnd_a, &rnd_b);
        assert_eq!(sv.len(), 32);
        assert_eq!(sv[0], 0xA5);
        assert_eq!(sv[1], 0x5A);
        assert_eq!(sv[2], 0x00);
        assert_eq!(sv[3], 0x01);
        assert_eq!(sv[4], 0x00);
        assert_eq!(sv[5], 0x80);
    }

    #[test]
    fn padding_round_trip() {
        let data = vec![1, 2, 3, 4, 5];
        let padded = pad_iso9797_m2(&data);
        assert_eq!(padded.len(), 16);
        assert_eq!(padded[5], 0x80);
        let unpadded = unpad_iso9797_m2(&padded).unwrap();
        assert_eq!(unpadded, data);
    }

    #[test]
    fn padding_block_aligned() {
        let data = vec![0xFFu8; 16];
        let padded = pad_iso9797_m2(&data);
        assert_eq!(padded.len(), 32); // full padding block added
        let unpadded = unpad_iso9797_m2(&padded).unwrap();
        assert_eq!(unpadded, data);
    }

    #[test]
    fn session_key_derivation_deterministic() {
        let key = [0x00u8; 16];
        let rnd_a = [0x01u8; 16];
        let rnd_b = [0x02u8; 16];
        let (ks_enc1, ks_mac1) = derive_session_keys(&key, &rnd_a, &rnd_b);
        let (ks_enc2, ks_mac2) = derive_session_keys(&key, &rnd_a, &rnd_b);
        assert_eq!(ks_enc1, ks_enc2);
        assert_eq!(ks_mac1, ks_mac2);
        assert_ne!(ks_enc1, ks_mac1); // different labels produce different keys
    }

    #[test]
    fn build_change_key_command_matches_an12196_case_1_example() {
        let mut session = Ev2Session {
            auth_key_no: 0x00,
            ti: hex::decode("7614281A").unwrap().try_into().unwrap(),
            cmd_ctr: 0x0002,
            ks_enc: hex::decode("4CF3CB41A22583A61E89B158D252FC53").unwrap().try_into().unwrap(),
            ks_mac: hex::decode("5529860B2FC5FB6154B7F28361D30BF9").unwrap().try_into().unwrap(),
        };
        let new_key: [u8; 16] = hex::decode("F3847D627727ED3BC9C4CC050489B966")
            .unwrap()
            .try_into()
            .unwrap();
        let old_key = [0u8; 16];
        let plaintext = ntag424::build_change_key_data_other(&new_key, &old_key, 0x02, 0x01);
        let apdu = session.build_change_key_command(&[0x02], &plaintext, false);

        let expected = hex::decode(
            "90C4000029022CF362B7BF4311FF3BE1DAA295E8C68DE09050560D19B9E16C2393AE9CD1FAC75D0CE20BCD1D06E600",
        )
        .unwrap();

        assert_eq!(apdu, expected);
    }
}
