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()
}