use aes_gcm::{
    aead::{Aead, KeyInit},
    Aes256Gcm, Nonce,
};
use anyhow::{Context, Result};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use rand::Rng;
use std::fs;
use std::path::Path;

const KEY_LEN: usize = 32;
const NONCE_LEN: usize = 12;
const SALT_LEN: usize = 32;

/// Encrypt data with password
pub fn encrypt(data: &[u8], password: &str) -> Result<String> {
    let mut salt = [0u8; SALT_LEN];
    rand::thread_rng().fill(&mut salt);
    let mut nonce_bytes = [0u8; NONCE_LEN];
    rand::thread_rng().fill(&mut nonce_bytes);
    
    let key = derive_key(password, &salt)?;
    let cipher = Aes256Gcm::new_from_slice(&key)
        .context("Failed to create cipher")?;
    let nonce = Nonce::from_slice(&nonce_bytes);
    
    let encrypted = cipher
        .encrypt(nonce, data)
        .map_err(|e| anyhow::anyhow!("Encryption failed: {:?}", e))?;
    
    // Combine salt + nonce + encrypted data
    let mut result = Vec::with_capacity(SALT_LEN + NONCE_LEN + encrypted.len());
    result.extend_from_slice(&salt);
    result.extend_from_slice(&nonce_bytes);
    result.extend_from_slice(&encrypted);
    
    Ok(BASE64.encode(&result))
}

/// Decrypt data with password
pub fn decrypt(encrypted_data: &str, password: &str) -> Result<Vec<u8>> {
    let data = BASE64.decode(encrypted_data)
        .context("Invalid base64 encoding")?;
    
    if data.len() < SALT_LEN + NONCE_LEN {
        anyhow::bail!("Invalid encrypted data format");
    }
    
    let salt = &data[..SALT_LEN];
    let nonce_bytes = &data[SALT_LEN..SALT_LEN + NONCE_LEN];
    let encrypted = &data[SALT_LEN + NONCE_LEN..];
    
    let key = derive_key(password, salt)?;
    let cipher = Aes256Gcm::new_from_slice(&key)
        .context("Failed to create cipher")?;
    let nonce = Nonce::from_slice(nonce_bytes);
    
    let decrypted = cipher
        .decrypt(nonce, encrypted)
        .map_err(|e| anyhow::anyhow!("Decryption failed: {:?}", e))?;
    
    Ok(decrypted)
}

/// Derive key from password using simple method
fn derive_key(password: &str, salt: &[u8]) -> Result<[u8; KEY_LEN]> {
    use sha2::{Sha256, Digest};
    
    let mut hasher = Sha256::new();
    hasher.update(salt);
    hasher.update(password.as_bytes());
    hasher.update(b"quicommit_key_derivation_v1");
    
    let hash = hasher.finalize();
    let mut key = [0u8; KEY_LEN];
    key.copy_from_slice(&hash[..KEY_LEN]);
    
    Ok(key)
}

/// Encrypt and save to file
pub fn encrypt_to_file(data: &[u8], password: &str, path: &Path) -> Result<()> {
    let encrypted = encrypt(data, password)?;
    fs::write(path, encrypted)
        .with_context(|| format!("Failed to write encrypted file: {:?}", path))?;
    Ok(())
}

/// Decrypt from file
pub fn decrypt_from_file(path: &Path, password: &str) -> Result<Vec<u8>> {
    let encrypted = fs::read_to_string(path)
        .with_context(|| format!("Failed to read encrypted file: {:?}", path))?;
    decrypt(&encrypted, password)
}

/// Generate a secure random token
pub fn generate_token(length: usize) -> String {
    const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
    let mut rng = rand::thread_rng();
    
    (0..length)
        .map(|_| {
            let idx = rng.gen_range(0..CHARSET.len());
            CHARSET[idx] as char
        })
        .collect()
}

/// Hash data using SHA-256
pub fn sha256(data: &[u8]) -> String {
    use sha2::{Digest, Sha256};
    let mut hasher = Sha256::new();
    hasher.update(data);
    hex::encode(hasher.finalize())
}

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

    #[test]
    fn test_encrypt_decrypt() {
        let data = b"Hello, World!";
        let password = "my_secret_password";
        
        let encrypted = encrypt(data, password).unwrap();
        let decrypted = decrypt(&encrypted, password).unwrap();
        
        assert_eq!(data.to_vec(), decrypted);
    }

    #[test]
    fn test_wrong_password() {
        let data = b"Hello, World!";
        let encrypted = encrypt(data, "correct_password").unwrap();
        
        assert!(decrypt(&encrypted, "wrong_password").is_err());
    }
}
