use aes_gcm::{
    aead::{Aead, KeyInit, OsRng},
    aead::rand_core::RngCore,
    Aes256Gcm, Key, Nonce,
};
use argon2::{Algorithm, Argon2, Params, Version};

// ---------------------------------------------------------------------------
// VaultService
// ---------------------------------------------------------------------------

/// Wraps a 256-bit symmetric key and exposes encrypt / decrypt helpers.
///
/// The wire format is:
///   `v1:{iv_hex}:{sealed_hex}`
///
/// where:
///   - `iv_hex`     — 12 bytes (96-bit) random nonce, hex-encoded (24 chars)
///   - `sealed_hex` — ciphertext + 16-byte GCM auth tag, hex-encoded
///
/// The version prefix allows a future migration to a different algorithm
/// without breaking existing stored blobs.
pub struct VaultService {
    key: [u8; 32],
}

impl VaultService {
    pub fn new(key: [u8; 32]) -> Self {
        Self { key }
    }

    /// Encrypt `plaintext` and return a `v1:{iv_hex}:{sealed_hex}` blob.
    pub fn encrypt(&self, plaintext: &str) -> Result<String, String> {
        // Build the AES-256-GCM cipher from our key.
        let key = Key::<Aes256Gcm>::from_slice(&self.key);
        let cipher = Aes256Gcm::new(key);

        // Generate a random 12-byte nonce (96-bit is the GCM standard).
        let mut iv_bytes = [0u8; 12];
        OsRng.fill_bytes(&mut iv_bytes);
        let nonce = Nonce::from_slice(&iv_bytes);

        // Encrypt. The returned Vec includes the 16-byte GCM auth tag appended.
        let sealed = cipher
            .encrypt(nonce, plaintext.as_bytes())
            .map_err(|e| format!("encryption failed: {e}"))?;

        Ok(format!("v1:{}:{}", hex::encode(iv_bytes), hex::encode(sealed)))
    }

    /// Decrypt a `v1:{iv_hex}:{sealed_hex}` blob and return the plaintext.
    pub fn decrypt(&self, blob: &str) -> Result<String, String> {
        // Parse the version-prefixed format.
        let parts: Vec<&str> = blob.splitn(3, ':').collect();
        if parts.len() != 3 || parts[0] != "v1" {
            return Err(format!(
                "unsupported vault blob format — expected 'v1:iv:sealed', got prefix '{}'",
                parts.first().copied().unwrap_or("<empty>")
            ));
        }

        let iv_bytes = hex::decode(parts[1])
            .map_err(|e| format!("invalid iv hex: {e}"))?;
        let sealed = hex::decode(parts[2])
            .map_err(|e| format!("invalid sealed hex: {e}"))?;

        if iv_bytes.len() != 12 {
            return Err(format!(
                "iv must be 12 bytes, got {}",
                iv_bytes.len()
            ));
        }

        let key = Key::<Aes256Gcm>::from_slice(&self.key);
        let cipher = Aes256Gcm::new(key);
        let nonce = Nonce::from_slice(&iv_bytes);

        let plaintext_bytes = cipher
            .decrypt(nonce, sealed.as_ref())
            .map_err(|e| format!("decryption failed (wrong key or corrupted data): {e}"))?;

        String::from_utf8(plaintext_bytes)
            .map_err(|e| format!("decrypted bytes are not valid UTF-8: {e}"))
    }
}

// ---------------------------------------------------------------------------
// Key derivation
// ---------------------------------------------------------------------------

/// Derive a 256-bit key from `password` + `salt` using Argon2id.
///
/// Parameters chosen to be strong enough for a desktop vault while completing
/// in well under a second on modern hardware:
///   t = 3 iterations
///   m = 65536 KiB (64 MiB) memory
///   p = 4 parallelism lanes
pub fn derive_key(password: &str, salt: &[u8]) -> [u8; 32] {
    let params = Params::new(
        65536, // m_cost: 64 MiB
        3,     // t_cost: iterations
        4,     // p_cost: parallelism
        Some(32), // output length
    )
    .expect("Argon2 params are hard-coded and always valid");

    let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);

    let mut output_key = [0u8; 32];
    argon2
        .hash_password_into(password.as_bytes(), salt, &mut output_key)
        .expect("Argon2id key derivation failed");

    output_key
}

/// Generate a cryptographically random 32-byte salt.
pub fn generate_salt() -> [u8; 32] {
    let mut salt = [0u8; 32];
    OsRng.fill_bytes(&mut salt);
    salt
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    fn make_vault() -> VaultService {
        let key = derive_key("test-password", &generate_salt());
        VaultService::new(key)
    }

    #[test]
    fn round_trip_ascii() {
        let vault = make_vault();
        let original = "hunter2";
        let blob = vault.encrypt(original).unwrap();
        assert!(blob.starts_with("v1:"), "blob must have v1 prefix");
        let recovered = vault.decrypt(&blob).unwrap();
        assert_eq!(recovered, original);
    }

    #[test]
    fn round_trip_unicode() {
        let vault = make_vault();
        let original = "p@ssw0rd — こんにちは 🔐";
        let recovered = vault.decrypt(&vault.encrypt(original).unwrap()).unwrap();
        assert_eq!(recovered, original);
    }

    #[test]
    fn different_nonces_per_encrypt() {
        let vault = make_vault();
        let b1 = vault.encrypt("same").unwrap();
        let b2 = vault.encrypt("same").unwrap();
        // Each call generates a fresh nonce so the blobs must differ.
        assert_ne!(b1, b2);
    }

    #[test]
    fn wrong_key_fails_decryption() {
        let vault1 = make_vault();
        let vault2 = make_vault(); // different salt → different key
        let blob = vault1.encrypt("secret").unwrap();
        assert!(vault2.decrypt(&blob).is_err());
    }

    #[test]
    fn malformed_blob_rejected() {
        let vault = make_vault();
        assert!(vault.decrypt("garbage").is_err());
        assert!(vault.decrypt("v2:aabb:ccdd").is_err());
        assert!(vault.decrypt("v1:not-hex:ccdd").is_err());
    }
}