use aes_gcm::{
    aead::{Aead, KeyInit, OsRng},
    Aes256Gcm, Nonce,
};
use anyhow::{Context, Result};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use rand::RngCore;
use sha2::{Digest, Sha256};

/// Derive a 256-bit key from a string using SHA-256.
fn derive_key(key: &str) -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(key.as_bytes());
    hasher.finalize().into()
}

/// Encrypt a plaintext string using AES-256-GCM.
///
/// Returns a base64-encoded string containing the nonce + ciphertext.
pub fn encrypt(plaintext: &str, key: &str) -> Result<String> {
    let key_bytes = derive_key(key);
    let cipher = Aes256Gcm::new_from_slice(&key_bytes)
        .context("Failed to create AES cipher")?;

    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.as_bytes())
        .map_err(|e| anyhow::anyhow!("Encryption failed: {e}"))?;

    // Prepend nonce to ciphertext
    let mut combined = Vec::with_capacity(12 + ciphertext.len());
    combined.extend_from_slice(&nonce_bytes);
    combined.extend_from_slice(&ciphertext);

    Ok(BASE64.encode(&combined))
}

/// Decrypt an AES-256-GCM encrypted string.
pub fn decrypt(encoded: &str, key: &str) -> Result<String> {
    let combined = BASE64.decode(encoded).context("Invalid base64")?;

    if combined.len() < 12 {
        anyhow::bail!("Ciphertext too short");
    }

    let (nonce_bytes, ciphertext) = combined.split_at(12);
    let key_bytes = derive_key(key);
    let cipher = Aes256Gcm::new_from_slice(&key_bytes)
        .context("Failed to create AES cipher")?;
    let nonce = Nonce::from_slice(nonce_bytes);

    let plaintext = cipher
        .decrypt(nonce, ciphertext)
        .map_err(|e| anyhow::anyhow!("Decryption failed: {e}"))?;

    String::from_utf8(plaintext).context("Decrypted data is not valid UTF-8")
}

/// Generate a cryptographically random hex token of the specified byte length.
pub fn generate_token(length: usize) -> String {
    let mut bytes = vec![0u8; length];
    OsRng.fill_bytes(&mut bytes);
    hex::encode(&bytes)
}
