use aes_gcm::{
    aead::{Aead, KeyInit, OsRng},
    Aes128Gcm, Nonce,
};
use aes_gcm::aead::rand_core::RngCore;
use std::fs;
use std::path::PathBuf;

pub fn validate_key(key: &str) -> Result<[u8; 16], String> {
    let key_bytes = key.as_bytes();
    if key_bytes.len() != 16 {
        return Err(format!("密钥长度无效: 需要16字节，当前为{}字节", key_bytes.len()));
    }
    let mut arr = [0u8; 16];
    arr.copy_from_slice(key_bytes);
    Ok(arr)
}

/// 加密字节数据，返回 nonce(12字节) + 密文
pub fn encrypt_bytes(plain_text: &[u8], key_bytes: &[u8; 16]) -> Result<Vec<u8>, String> {
    let cipher = Aes128Gcm::new_from_slice(key_bytes)
        .map_err(|e| format!("创建 cipher 失败: {}", e))?;

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

    let cipher_text = cipher
        .encrypt(nonce, plain_text)
        .map_err(|e| format!("加密失败: {}", e))?;

    let mut result = Vec::with_capacity(12 + cipher_text.len());
    result.extend_from_slice(&nonce_bytes);
    result.extend_from_slice(&cipher_text);
    Ok(result)
}

/// 解密字节数据，输入格式为 nonce(12字节) + 密文
pub fn decrypt_bytes(data: &[u8], key_bytes: &[u8; 16]) -> Result<Vec<u8>, String> {
    if data.len() < 12 {
        return Err("密文长度无效：小于 nonce 长度（12字节）".to_string());
    }

    let cipher = Aes128Gcm::new_from_slice(key_bytes)
        .map_err(|e| format!("创建 cipher 失败: {}", e))?;

    let (nonce_bytes, cipher_text) = data.split_at(12);
    let nonce = Nonce::from_slice(nonce_bytes);

    cipher
        .decrypt(nonce, cipher_text)
        .map_err(|_| "解密失败：密钥错误或文件已损坏".to_string())
}

pub fn encrypt_file(input: &PathBuf, output: &PathBuf, key: &str) -> Result<(), String> {
    let plain_text = fs::read(input).map_err(|e| format!("读取输入文件失败: {}", e))?;
    let key_bytes = validate_key(key.trim())?;
    let result = encrypt_bytes(&plain_text, &key_bytes)?;

    write_file(output, &result)?;
    println!("加密成功: {:?} -> {:?}", input, output);
    Ok(())
}

pub fn decrypt_file(input: &PathBuf, output: &PathBuf, key: &str) -> Result<(), String> {
    let data = fs::read(input).map_err(|e| format!("读取输入文件失败: {}", e))?;
    let key_bytes = validate_key(key.trim())?;
    let plain_text = decrypt_bytes(&data, &key_bytes)?;

    write_file(output, &plain_text)?;
    println!("解密成功: {:?} -> {:?}", input, output);
    Ok(())
}

fn write_file(path: &PathBuf, data: &[u8]) -> Result<(), String> {
    if let Some(parent) = path.parent() {
        if !parent.as_os_str().is_empty() {
            fs::create_dir_all(parent).map_err(|e| format!("创建输出目录失败: {}", e))?;
        }
    }
    fs::write(path, data).map_err(|e| format!("写入输出文件失败: {}", e))
}
