//! 加密工具

use crate::{Error, Result};
use ring::aead::{Aad, BoundKey, Nonce, NonceSequence, SealingKey, OpeningKey, UnboundKey, AES_256_GCM};
use ring::rand::{SecureRandom, SystemRandom};

/// 密钥长度
pub const KEY_SIZE: usize = 32;
/// Nonce长度
pub const NONCE_SIZE: usize = 12;
/// 认证标签长度
pub const TAG_SIZE: usize = 16;

/// 生成随机密钥
pub fn generate_key() -> [u8; KEY_SIZE] {
    let rng = SystemRandom::new();
    let mut key = [0u8; KEY_SIZE];
    rng.fill(&mut key).expect("Failed to generate key");
    key
}

/// 生成随机Nonce
pub fn generate_nonce() -> [u8; NONCE_SIZE] {
    let rng = SystemRandom::new();
    let mut nonce = [0u8; NONCE_SIZE];
    rng.fill(&mut nonce).expect("Failed to generate nonce");
    nonce
}

/// 计数器Nonce序列
struct CounterNonceSequence {
    counter: u64,
    prefix: [u8; 4],
}

impl CounterNonceSequence {
    fn new() -> Self {
        let rng = SystemRandom::new();
        let mut prefix = [0u8; 4];
        rng.fill(&mut prefix).expect("Failed to generate prefix");
        Self { counter: 0, prefix }
    }
}

impl NonceSequence for CounterNonceSequence {
    fn advance(&mut self) -> std::result::Result<Nonce, ring::error::Unspecified> {
        let mut nonce = [0u8; NONCE_SIZE];
        nonce[..4].copy_from_slice(&self.prefix);
        nonce[4..].copy_from_slice(&self.counter.to_be_bytes());
        self.counter += 1;
        Nonce::try_assume_unique_for_key(&nonce)
    }
}

/// 加密器
pub struct Encryptor {
    sealing_key: SealingKey<CounterNonceSequence>,
}

impl Encryptor {
    /// 创建新的加密器
    pub fn new(key: &[u8; KEY_SIZE]) -> Result<Self> {
        let unbound_key = UnboundKey::new(&AES_256_GCM, key)
            .map_err(|_| Error::CryptoError("Invalid key".into()))?;
        let sealing_key = SealingKey::new(unbound_key, CounterNonceSequence::new());
        Ok(Self { sealing_key })
    }

    /// 加密数据
    pub fn encrypt(&mut self, plaintext: &[u8]) -> Result<Vec<u8>> {
        let mut in_out = plaintext.to_vec();
        // 预留TAG空间
        in_out.extend_from_slice(&[0u8; TAG_SIZE]);
        
        self.sealing_key
            .seal_in_place_append_tag(Aad::empty(), &mut in_out)
            .map_err(|_| Error::CryptoError("Encryption failed".into()))?;
        
        Ok(in_out)
    }
}

/// 解密器
pub struct Decryptor {
    opening_key: OpeningKey<CounterNonceSequence>,
}

impl Decryptor {
    /// 创建新的解密器
    pub fn new(key: &[u8; KEY_SIZE]) -> Result<Self> {
        let unbound_key = UnboundKey::new(&AES_256_GCM, key)
            .map_err(|_| Error::CryptoError("Invalid key".into()))?;
        let opening_key = OpeningKey::new(unbound_key, CounterNonceSequence::new());
        Ok(Self { opening_key })
    }

    /// 解密数据
    pub fn decrypt(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>> {
        let mut in_out = ciphertext.to_vec();
        
        let plaintext = self.opening_key
            .open_in_place(Aad::empty(), &mut in_out)
            .map_err(|_| Error::CryptoError("Decryption failed".into()))?;
        
        Ok(plaintext.to_vec())
    }
}

/// 密码哈希
pub mod password {
    use super::*;
    use ring::pbkdf2;
    use std::num::NonZeroU32;

    const CREDENTIAL_LEN: usize = 32;
    const ITERATIONS: u32 = 100_000;

    /// 哈希密码
    pub fn hash_password(password: &str, salt: &[u8]) -> [u8; CREDENTIAL_LEN] {
        let mut hash = [0u8; CREDENTIAL_LEN];
        pbkdf2::derive(
            pbkdf2::PBKDF2_HMAC_SHA256,
            NonZeroU32::new(ITERATIONS).unwrap(),
            salt,
            password.as_bytes(),
            &mut hash,
        );
        hash
    }

    /// 验证密码
    pub fn verify_password(password: &str, salt: &[u8], hash: &[u8]) -> bool {
        pbkdf2::verify(
            pbkdf2::PBKDF2_HMAC_SHA256,
            NonZeroU32::new(ITERATIONS).unwrap(),
            salt,
            password.as_bytes(),
            hash,
        )
        .is_ok()
    }

    /// 生成盐值
    pub fn generate_salt() -> [u8; 16] {
        let rng = SystemRandom::new();
        let mut salt = [0u8; 16];
        rng.fill(&mut salt).expect("Failed to generate salt");
        salt
    }
}

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

    #[test]
    fn test_encrypt_decrypt() {
        let key = generate_key();
        let mut encryptor = Encryptor::new(&key).unwrap();
        let mut decryptor = Decryptor::new(&key).unwrap();

        let plaintext = b"Hello, World!";
        let ciphertext = encryptor.encrypt(plaintext).unwrap();
        let decrypted = decryptor.decrypt(&ciphertext).unwrap();

        assert_eq!(plaintext.as_slice(), decrypted.as_slice());
    }

    #[test]
    fn test_password_hash() {
        let password = "my_secure_password";
        let salt = password::generate_salt();
        let hash = password::hash_password(password, &salt);

        assert!(password::verify_password(password, &salt, &hash));
        assert!(!password::verify_password("wrong_password", &salt, &hash));
    }
}
