//! age encryption wrappers for team vault use.
//!
//! Provides `encrypt` and `decrypt` over age's X25519 recipient model so
//! team vaults can be encrypted to multiple recipients without a shared password.

use std::io::{Read, Write};

use crate::errors::{SafeError, SafeResult};

/// Encrypt data to one or more age recipients (X25519 public keys).
pub fn encrypt_to_recipients(
    recipients: &[Box<dyn age::Recipient + Send>],
    plaintext: &[u8],
) -> SafeResult<Vec<u8>> {
    let refs: Vec<&dyn age::Recipient> = recipients
        .iter()
        .map(|r| r.as_ref() as &dyn age::Recipient)
        .collect();
    let encryptor =
        age::Encryptor::with_recipients(refs.into_iter()).map_err(|e| SafeError::Crypto {
            context: format!("age encrypt init: {e}"),
        })?;
    let mut encrypted = Vec::new();
    let mut writer = encryptor
        .wrap_output(&mut encrypted)
        .map_err(|e| SafeError::Crypto {
            context: format!("age encrypt: {e}"),
        })?;
    writer.write_all(plaintext).map_err(|e| SafeError::Crypto {
        context: format!("age write: {e}"),
    })?;
    writer.finish().map_err(|e| SafeError::Crypto {
        context: format!("age finish: {e}"),
    })?;
    Ok(encrypted)
}

/// Decrypt data using one or more age identities (X25519 secret keys).
pub fn decrypt_with_identities(
    identities: &[Box<dyn age::Identity>],
    ciphertext: &[u8],
) -> SafeResult<Vec<u8>> {
    let decryptor = age::Decryptor::new(ciphertext).map_err(|e| SafeError::Crypto {
        context: format!("age decryptor init: {e}"),
    })?;
    let mut reader = decryptor
        .decrypt(identities.iter().map(|i| i.as_ref()))
        .map_err(|e| SafeError::Crypto {
            context: format!("age decrypt: {e}"),
        })?;
    let mut plaintext = Vec::new();
    reader
        .read_to_end(&mut plaintext)
        .map_err(|e| SafeError::Crypto {
            context: format!("age read: {e}"),
        })?;
    Ok(plaintext)
}

/// Parse an age identity file (e.g. `~/.age/key.txt`).
pub fn load_identities(path: &std::path::Path) -> SafeResult<Vec<Box<dyn age::Identity>>> {
    let content = std::fs::read_to_string(path).map_err(|e| SafeError::Crypto {
        context: format!("read identity file {}: {e}", path.display()),
    })?;
    let file =
        age::IdentityFile::from_buffer(content.as_bytes()).map_err(|e| SafeError::Crypto {
            context: format!("parse identity file: {e}"),
        })?;
    let identities = file.into_identities().map_err(|e| SafeError::Crypto {
        context: format!("extract identities: {e}"),
    })?;
    Ok(identities)
}

/// Generate a new age X25519 identity (keypair).
/// Returns `(identity_string, recipient_string)`.
pub fn generate_identity() -> (String, String) {
    use age::secrecy::ExposeSecret;
    let identity = age::x25519::Identity::generate();
    let recipient = identity.to_public().to_string();
    let secret = identity.to_string();
    (secret.expose_secret().to_string(), recipient)
}

/// Parse an age recipient public key string (e.g. `age1...`).
pub fn parse_recipient(pubkey: &str) -> SafeResult<Box<dyn age::Recipient + Send>> {
    let r: age::x25519::Recipient = pubkey.parse().map_err(|e: &str| SafeError::Crypto {
        context: format!("invalid age recipient '{pubkey}': {e}"),
    })?;
    Ok(Box::new(r))
}

/// Parse multiple recipient public key strings.
pub fn parse_recipients(pubkeys: &[String]) -> SafeResult<Vec<Box<dyn age::Recipient + Send>>> {
    pubkeys.iter().map(|pk| parse_recipient(pk)).collect()
}

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

    #[test]
    fn encrypt_decrypt_roundtrip() {
        let (secret, pubkey) = generate_identity();

        let recipients = vec![parse_recipient(&pubkey).unwrap()];
        let plaintext = b"hello, age!";
        let ciphertext = encrypt_to_recipients(&recipients, plaintext).unwrap();

        let identities = age::IdentityFile::from_buffer(secret.as_bytes())
            .unwrap()
            .into_identities()
            .unwrap();
        let decrypted = decrypt_with_identities(&identities, &ciphertext).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn multi_recipient_both_can_decrypt() {
        let (secret1, pubkey1) = generate_identity();
        let (secret2, pubkey2) = generate_identity();

        let recipients = vec![
            parse_recipient(&pubkey1).unwrap(),
            parse_recipient(&pubkey2).unwrap(),
        ];
        let ciphertext = encrypt_to_recipients(&recipients, b"shared secret").unwrap();

        let ids1 = age::IdentityFile::from_buffer(secret1.as_bytes())
            .unwrap()
            .into_identities()
            .unwrap();
        assert_eq!(
            decrypt_with_identities(&ids1, &ciphertext).unwrap(),
            b"shared secret"
        );

        let ids2 = age::IdentityFile::from_buffer(secret2.as_bytes())
            .unwrap()
            .into_identities()
            .unwrap();
        assert_eq!(
            decrypt_with_identities(&ids2, &ciphertext).unwrap(),
            b"shared secret"
        );
    }

    #[test]
    fn wrong_identity_fails() {
        let (_secret1, pubkey1) = generate_identity();
        let (secret2, _pubkey2) = generate_identity();

        let recipients = vec![parse_recipient(&pubkey1).unwrap()];
        let ciphertext = encrypt_to_recipients(&recipients, b"secret").unwrap();

        let ids2 = age::IdentityFile::from_buffer(secret2.as_bytes())
            .unwrap()
            .into_identities()
            .unwrap();
        assert!(decrypt_with_identities(&ids2, &ciphertext).is_err());
    }
}
