//! AES-256-GCM，与 pythonbackend `app/core/aes_crypto.py` 一致。

use rand::RngCore;

use aes_gcm::{
    aead::{Aead, KeyInit},
    Aes256Gcm, Nonce,
};
use base64::Engine;
use serde_json::Value;

pub const ENCRYPTED_HEADER: &str = "X-AICLIENT-Encrypted";
const ENVELOPE_VERSION: i64 = 1;

pub fn parse_aes_key(raw: &str) -> Option<[u8; 32]> {
    let s = raw.trim();
    if s.is_empty() {
        return None;
    }
    let bytes = if s.len() == 64 && s.chars().all(|c| c.is_ascii_hexdigit()) {
        hex::decode(s).ok()?
    } else {
        base64::engine::general_purpose::STANDARD
            .decode(s)
            .ok()?
    };
    if bytes.len() != 32 {
        return None;
    }
    let mut key = [0u8; 32];
    key.copy_from_slice(&bytes);
    Some(key)
}

pub fn aes_key_from_env() -> Option<[u8; 32]> {
    let s="98cde5e64ee0da6ed0c5842f43a950aa0d8ae5087c43c8d8ee5d10bd8d571bb4";
    parse_aes_key(&s)
}

pub fn encrypt_bytes(plaintext: &[u8], key: &[u8; 32]) -> Result<String, String> {
    let cipher = Aes256Gcm::new_from_slice(key).map_err(|e| e.to_string())?;
    let mut nonce_bytes = [0u8; 12];
    rand::RngCore::fill_bytes(&mut rand::thread_rng(), &mut nonce_bytes);
    let nonce = Nonce::from_slice(&nonce_bytes);
    let ciphertext = cipher
        .encrypt(nonce, plaintext)
        .map_err(|e| format!("加密失败: {e}"))?;
    let mut combined = nonce_bytes.to_vec();
    combined.extend(ciphertext);
    Ok(base64::engine::general_purpose::STANDARD.encode(&combined))
}

pub fn decrypt_payload(b64_payload: &str, key: &[u8; 32]) -> Result<Vec<u8>, String> {
    let raw = base64::engine::general_purpose::STANDARD
        .decode(b64_payload.trim())
        .map_err(|e| format!("Base64 解码失败: {e}"))?;
    if raw.len() < 13 {
        return Err("密文过短".into());
    }
    let (nonce_bytes, ciphertext) = raw.split_at(12);
    let cipher = Aes256Gcm::new_from_slice(key).map_err(|e| e.to_string())?;
    let nonce = Nonce::from_slice(nonce_bytes);
    cipher
        .decrypt(nonce, ciphertext)
        .map_err(|e| format!("解密失败: {e}"))
}

pub fn wrap_encrypted(plaintext: &[u8], key: &[u8; 32]) -> Result<String, String> {
    let payload = encrypt_bytes(plaintext, key)?;
    let envelope = serde_json::json!({
        "v": ENVELOPE_VERSION,
        "payload": payload,
    });
    serde_json::to_string(&envelope).map_err(|e| e.to_string())
}

pub fn unwrap_encrypted(body: &str, key: &[u8; 32]) -> Result<Vec<u8>, String> {
    let v: Value = serde_json::from_str(body).map_err(|e| format!("信封 JSON 无效: {e}"))?;
    let version = v.get("v").and_then(|x| x.as_i64());
    let payload = v
        .get("payload")
        .and_then(|x| x.as_str())
        .ok_or_else(|| "缺少 payload".to_string())?;
    if version != Some(ENVELOPE_VERSION) {
        return Err("不支持的加密版本".into());
    }
    decrypt_payload(payload, key)
}

pub fn crypto_enabled() -> bool {
    aes_key_from_env().is_some()
}