use aes_gcm::{
    Aes256Gcm, Key, Nonce,
    aead::{Aead, AeadCore, KeyInit, OsRng},
};
use anyhow::{Context, Result, bail};
use serde_json::Value;

const NONCE_LEN: usize = 12;

// ─── AES-256-GCM encrypt / decrypt ───────────────────────────────────────────

/// Encrypt plaintext bytes with AES-256-GCM.
/// Returns `nonce (12 B) || ciphertext+tag`.
pub fn encrypt(master_key: &[u8; 32], plaintext: &[u8]) -> Result<Vec<u8>> {
    let key = Key::<Aes256Gcm>::from_slice(master_key);
    let cipher = Aes256Gcm::new(key);
    let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
    let ciphertext = cipher
        .encrypt(&nonce, plaintext)
        .map_err(|e| anyhow::anyhow!("AES-256-GCM encryption failed: {}", e))?;
    let mut out = Vec::with_capacity(NONCE_LEN + ciphertext.len());
    out.extend_from_slice(&nonce);
    out.extend_from_slice(&ciphertext);
    Ok(out)
}

/// Decrypt `nonce (12 B) || ciphertext+tag` with AES-256-GCM.
pub fn decrypt(master_key: &[u8; 32], data: &[u8]) -> Result<Vec<u8>> {
    if data.len() < NONCE_LEN {
        bail!(
            "encrypted data too short ({}B); possibly corrupted",
            data.len()
        );
    }
    let (nonce_bytes, ciphertext) = data.split_at(NONCE_LEN);
    let key = Key::<Aes256Gcm>::from_slice(master_key);
    let cipher = Aes256Gcm::new(key);
    let nonce = Nonce::from_slice(nonce_bytes);
    cipher
        .decrypt(nonce, ciphertext)
        .map_err(|_| anyhow::anyhow!("decryption failed — wrong master key or corrupted data"))
}

// ─── JSON helpers ─────────────────────────────────────────────────────────────

/// Serialize a JSON Value and encrypt it. Returns the encrypted blob.
pub fn encrypt_json(master_key: &[u8; 32], value: &Value) -> Result<Vec<u8>> {
    let bytes = serde_json::to_vec(value).context("serialize JSON for encryption")?;
    encrypt(master_key, &bytes)
}

/// Decrypt an encrypted blob and deserialize it as a JSON Value.
pub fn decrypt_json(master_key: &[u8; 32], data: &[u8]) -> Result<Value> {
    let bytes = decrypt(master_key, data)?;
    serde_json::from_slice(&bytes).context("deserialize decrypted JSON")
}

// ─── Per-user key management (DEPRECATED — kept only for migration) ───────────

/// Generate a new random 32-byte per-user encryption key.
#[allow(dead_code)]
pub fn generate_user_key() -> [u8; 32] {
    use aes_gcm::aead::rand_core::RngCore;
    let mut key = [0u8; 32];
    OsRng.fill_bytes(&mut key);
    key
}

/// Wrap a per-user key with the server master key using AES-256-GCM.
#[allow(dead_code)]
pub fn wrap_user_key(server_master_key: &[u8; 32], user_key: &[u8; 32]) -> Result<Vec<u8>> {
    encrypt(server_master_key, user_key.as_ref())
}

/// Unwrap a per-user key using the server master key.
#[allow(dead_code)]
pub fn unwrap_user_key(server_master_key: &[u8; 32], wrapped: &[u8]) -> Result<[u8; 32]> {
    let bytes = decrypt(server_master_key, wrapped)?;
    if bytes.len() != 32 {
        bail!("unwrapped user key has unexpected length {}", bytes.len());
    }
    let mut key = [0u8; 32];
    key.copy_from_slice(&bytes);
    Ok(key)
}

// ─── Client-supplied key extraction ──────────────────────────────────────────

/// Parse a 64-char hex string (from X-Encryption-Key header) into a 32-byte key.
pub fn extract_key_from_hex(hex_str: &str) -> Result<[u8; 32]> {
    let bytes = hex::decode_hex(hex_str.trim())?;
    if bytes.len() != 32 {
        bail!(
            "X-Encryption-Key must be 64 hex chars (32 bytes), got {} bytes",
            bytes.len()
        );
    }
    let mut key = [0u8; 32];
    key.copy_from_slice(&bytes);
    Ok(key)
}

// ─── Server master key ────────────────────────────────────────────────────────

/// Load the server master key from `SERVER_MASTER_KEY` environment variable (64 hex chars).
pub fn load_master_key_auto() -> Result<[u8; 32]> {
    let hex_str = std::env::var("SERVER_MASTER_KEY").map_err(|_| {
        anyhow::anyhow!(
            "SERVER_MASTER_KEY is not set. \
             Generate one with: openssl rand -hex 32"
        )
    })?;

    if hex_str.is_empty() {
        bail!("SERVER_MASTER_KEY is set but empty");
    }

    let bytes = hex::decode_hex(hex_str.trim())?;
    if bytes.len() != 32 {
        bail!(
            "SERVER_MASTER_KEY must be 64 hex chars (32 bytes), got {} bytes",
            bytes.len()
        );
    }
    let mut key = [0u8; 32];
    key.copy_from_slice(&bytes);
    Ok(key)
}

// ─── Public hex helpers ───────────────────────────────────────────────────────

pub mod hex {
    use anyhow::{Result, bail};

    pub fn encode_hex(bytes: &[u8]) -> String {
        bytes.iter().map(|b| format!("{:02x}", b)).collect()
    }

    pub fn decode_hex(s: &str) -> Result<Vec<u8>> {
        let s = s.trim();
        if !s.len().is_multiple_of(2) {
            bail!("hex string has odd length");
        }
        (0..s.len())
            .step_by(2)
            .map(|i| u8::from_str_radix(&s[i..i + 2], 16).map_err(|e| anyhow::anyhow!("{}", e)))
            .collect()
    }
}

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

    #[test]
    fn roundtrip_encrypt_decrypt() {
        let key = [0x42u8; 32];
        let plaintext = b"hello world";
        let enc = encrypt(&key, plaintext).unwrap();
        let dec = decrypt(&key, &enc).unwrap();
        assert_eq!(dec, plaintext);
    }

    #[test]
    fn encrypt_produces_different_ciphertexts() {
        let key = [0x42u8; 32];
        let plaintext = b"hello world";
        let enc1 = encrypt(&key, plaintext).unwrap();
        let enc2 = encrypt(&key, plaintext).unwrap();
        assert_ne!(enc1, enc2);
    }

    #[test]
    fn wrong_key_fails_decryption() {
        let key1 = [0x42u8; 32];
        let key2 = [0x43u8; 32];
        let enc = encrypt(&key1, b"secret").unwrap();
        assert!(decrypt(&key2, &enc).is_err());
    }

    #[test]
    fn json_roundtrip() {
        let key = [0x42u8; 32];
        let value = serde_json::json!({"token": "abc123", "password": "hunter2"});
        let enc = encrypt_json(&key, &value).unwrap();
        let dec = decrypt_json(&key, &enc).unwrap();
        assert_eq!(dec, value);
    }

    #[test]
    fn user_key_wrap_unwrap_roundtrip() {
        let server_key = [0xABu8; 32];
        let user_key = [0xCDu8; 32];
        let wrapped = wrap_user_key(&server_key, &user_key).unwrap();
        let unwrapped = unwrap_user_key(&server_key, &wrapped).unwrap();
        assert_eq!(unwrapped, user_key);
    }

    #[test]
    fn user_key_wrap_wrong_server_key_fails() {
        let server_key1 = [0xABu8; 32];
        let server_key2 = [0xEFu8; 32];
        let user_key = [0xCDu8; 32];
        let wrapped = wrap_user_key(&server_key1, &user_key).unwrap();
        assert!(unwrap_user_key(&server_key2, &wrapped).is_err());
    }
}