use aes_gcm::{
    aead::{Aead, KeyInit, OsRng},
    Aes256Gcm, Nonce,
};
use argon2::{Argon2, PasswordHasher, password_hash::SaltString};
use argon2::password_hash::rand_core::RngCore;

use super::error::{ClosureError, Result};

/// Ciphertext format: salt(16) || nonce(12) || ciphertext || tag(16)
pub fn encrypt(plaintext: &[u8], password: &str) -> Result<Vec<u8>> {
    let salt = SaltString::generate(&mut OsRng);
    let argon2 = Argon2::default();
    let password_hash = argon2
        .hash_password(password.as_bytes(), &salt)
        .map_err(|e| ClosureError::Crypto(e.to_string()))?;
    let key_bytes = password_hash.hash.ok_or_else(|| {
        ClosureError::Crypto("argon2 hash missing".to_string())
    })?;

    let key = aes_gcm::Key::<Aes256Gcm>::from_slice(key_bytes.as_ref());
    let cipher = Aes256Gcm::new(key);

    let mut nonce_bytes = [0u8; 12];
    OsRng.fill_bytes(&mut nonce_bytes);
    let nonce = Nonce::from_slice(&nonce_bytes);

    let ciphertext = cipher
        .encrypt(nonce, plaintext)
        .map_err(|e| ClosureError::Crypto(e.to_string()))?;

    let salt_bytes = salt.as_str().as_bytes();
    let salt_len = salt_bytes.len() as u8;

    let mut result = Vec::with_capacity(1 + salt_bytes.len() + nonce_bytes.len() + ciphertext.len());
    result.push(salt_len);
    result.extend_from_slice(salt_bytes);
    result.extend_from_slice(&nonce_bytes);
    result.extend_from_slice(&ciphertext);

    Ok(result)
}

pub fn decrypt(ciphertext: &[u8], password: &str) -> Result<Vec<u8>> {
    if ciphertext.len() < 1 {
        return Err(ClosureError::Crypto("ciphertext too short".to_string()));
    }

    let salt_len = ciphertext[0] as usize;
    if ciphertext.len() < 1 + salt_len + 12 {
        return Err(ClosureError::Crypto("ciphertext too short".to_string()));
    }

    let salt_bytes = &ciphertext[1..1 + salt_len];
    let salt_str = std::str::from_utf8(salt_bytes)
        .map_err(|e| ClosureError::Crypto(format!("invalid salt: {}", e)))?;
    let salt = SaltString::from_b64(salt_str)
        .map_err(|e| ClosureError::Crypto(format!("invalid salt: {}", e)))?;

    let nonce_bytes = &ciphertext[1 + salt_len..1 + salt_len + 12];
    let encrypted = &ciphertext[1 + salt_len + 12..];

    let argon2 = Argon2::default();
    let password_hash = argon2
        .hash_password(password.as_bytes(), &salt)
        .map_err(|e| ClosureError::Crypto(e.to_string()))?;
    let key_bytes = password_hash.hash.ok_or_else(|| {
        ClosureError::Crypto("argon2 hash missing".to_string())
    })?;

    let key = aes_gcm::Key::<Aes256Gcm>::from_slice(key_bytes.as_ref());
    let cipher = Aes256Gcm::new(key);
    let nonce = Nonce::from_slice(nonce_bytes);

    let plaintext = cipher
        .decrypt(nonce, encrypted)
        .map_err(|_| ClosureError::InvalidPassword)?;

    Ok(plaintext)
}
