use aes::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use hmac::Mac;
type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;
type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;
type HmacSha256 = hmac::Hmac<sha2::Sha256>;

/// Generate a random 16-byte initialization vector using a CSPRNG.
pub fn generate_iv() -> [u8; 16] {
    let mut iv = [0u8; 16];
    rand::Fill::fill(&mut iv, &mut rand::rng());
    iv
}

/// Encrypt plaintext with AES-256-CBC and PKCS7 padding.
///
/// Returns ciphertext of size `((plaintext.len() / 16) + 1) * 16`.
/// PKCS7 always adds at least 1 byte of padding.
pub fn encrypt_data(plaintext: &[u8], key: &[u8; 32], iv: &[u8; 16]) -> Vec<u8> {
    let encrypted_size = ((plaintext.len() / 16) + 1) * 16;
    let mut buf = vec![0u8; encrypted_size];
    buf[..plaintext.len()].copy_from_slice(plaintext);

    let ct = Aes256CbcEnc::new(key.into(), iv.into())
        .encrypt_padded_mut::<Pkcs7>(&mut buf, plaintext.len())
        .expect("encryption buffer too small");

    // ct is a slice into buf of length encrypted_size
    ct.to_vec()
}

/// Decrypt ciphertext with AES-256-CBC and remove PKCS7 padding.
///
/// Returns the original plaintext data.
pub fn decrypt_data(ciphertext: &[u8], key: &[u8; 32], iv: &[u8; 16]) -> anyhow::Result<Vec<u8>> {
    let mut buf = ciphertext.to_vec();

    let pt = Aes256CbcDec::new(key.into(), iv.into())
        .decrypt_padded_mut::<Pkcs7>(&mut buf)
        .map_err(|_| anyhow::anyhow!("Decryption failed: invalid padding or wrong key"))?;

    Ok(pt.to_vec())
}

/// Compute HMAC-SHA-256 over IV || ciphertext.
///
/// HMAC input = IV (16 bytes) || ciphertext (encrypted_size bytes).
/// Returns 32-byte HMAC tag.
pub fn compute_hmac(key: &[u8; 32], iv: &[u8; 16], ciphertext: &[u8]) -> [u8; 32] {
    let mut mac =
        HmacSha256::new_from_slice(key).expect("HMAC can take key of any size");
    mac.update(iv);
    mac.update(ciphertext);
    mac.finalize().into_bytes().into()
}

/// Verify HMAC-SHA-256 over IV || ciphertext using constant-time comparison.
///
/// Returns true if the computed HMAC matches the expected value.
pub fn verify_hmac(
    key: &[u8; 32],
    iv: &[u8; 16],
    ciphertext: &[u8],
    expected: &[u8; 32],
) -> bool {
    let mut mac =
        HmacSha256::new_from_slice(key).expect("HMAC can take key of any size");
    mac.update(iv);
    mac.update(ciphertext);
    mac.verify_slice(expected).is_ok()
}

/// Compute SHA-256 hash of data.
///
/// Returns 32-byte digest. Used for integrity verification of original file content.
pub fn sha256_hash(data: &[u8]) -> [u8; 32] {
    use sha2::Digest;
    sha2::Sha256::digest(data).into()
}

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

    /// Test key matching legacy hardcoded value
    const TEST_KEY: [u8; 32] = [
        0x7A, 0x35, 0xC1, 0xD9, 0x4F, 0xE8, 0x2B, 0x6A,
        0x91, 0x0D, 0xF3, 0x58, 0xBC, 0x74, 0xA6, 0x1E,
        0x42, 0x8F, 0xD0, 0x63, 0xE5, 0x17, 0x9B, 0x2C,
        0xFA, 0x84, 0x06, 0xCD, 0x3E, 0x79, 0xB5, 0x50,
    ];

    #[test]
    fn test_encrypt_decrypt_roundtrip() {
        let plaintext = b"Hello, World!";
        let iv = [0u8; 16];
        let ciphertext = encrypt_data(plaintext, &TEST_KEY, &iv);
        let decrypted = decrypt_data(&ciphertext, &TEST_KEY, &iv).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_encrypt_decrypt_empty() {
        let plaintext = b"";
        let iv = [0u8; 16];
        let ciphertext = encrypt_data(plaintext, &TEST_KEY, &iv);
        let decrypted = decrypt_data(&ciphertext, &TEST_KEY, &iv).unwrap();
        assert_eq!(decrypted, plaintext.as_slice());
    }

    #[test]
    fn test_encrypted_size_formula() {
        let iv = [0u8; 16];
        // 5 bytes -> ((5/16)+1)*16 = 16
        assert_eq!(encrypt_data(b"Hello", &TEST_KEY, &iv).len(), 16);
        // 16 bytes -> ((16/16)+1)*16 = 32 (full padding block)
        assert_eq!(encrypt_data(&[0u8; 16], &TEST_KEY, &iv).len(), 32);
        // 0 bytes -> ((0/16)+1)*16 = 16
        assert_eq!(encrypt_data(b"", &TEST_KEY, &iv).len(), 16);
    }

    #[test]
    fn test_hmac_compute_verify() {
        let iv = [0xAA; 16];
        let ciphertext = b"some ciphertext data here!!12345";
        let hmac_tag = compute_hmac(&TEST_KEY, &iv, ciphertext);
        // Verify with correct tag
        assert!(verify_hmac(&TEST_KEY, &iv, ciphertext, &hmac_tag));
        // Verify with wrong tag
        let wrong_tag = [0u8; 32];
        assert!(!verify_hmac(&TEST_KEY, &iv, ciphertext, &wrong_tag));
    }

    #[test]
    fn test_sha256_known_value() {
        // SHA-256("Hello") from FORMAT.md Section 12.3
        let expected = hex!("185f8db32271fe25f561a6fc938b2e264306ec304eda518007d1764826381969");
        let result = sha256_hash(b"Hello");
        assert_eq!(result, expected);
    }

    #[test]
    fn test_sha256_empty() {
        let expected = hex!("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855");
        let result = sha256_hash(b"");
        assert_eq!(result, expected);
    }
}