use aes_gcm::aead::{Aead, KeyInit, OsRng};
use aes_gcm::{Aes256Gcm, Nonce};
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use base64::Engine;
use hkdf::Hkdf;
use rsa::signature::Verifier;
use rsa::{Pkcs1v15Sign, RsaPublicKey};
use sha2::{Digest, Sha256};

use crate::error::LicenseError;

const EMBEDDED_AES_SALT: &[u8] = b"craftlabs-license-salt-v1-2026-05";
const AES_NONCE_SIZE: usize = 12;

pub fn verify_rsa_signature(
    public_key_pem: &str,
    data: &[u8],
    signature_b64: &str,
) -> Result<(), LicenseError> {
    use rsa::pkcs8::DecodePublicKey;

    let pub_key = RsaPublicKey::from_public_key_pem(public_key_pem)
        .map_err(|_| LicenseError::InvalidSignature)?;

    let sig_bytes = URL_SAFE_NO_PAD
        .decode(signature_b64)
        .map_err(|_| LicenseError::InvalidFormat("signature base64"))?;

    let mut hasher = Sha256::new();
    hasher.update(data);
    let digest = hasher.finalize();

    pub_key
        .verify(Pkcs1v15Sign::new::<Sha256>(), &digest, sig_bytes.as_slice())
        .map_err(|_| LicenseError::SignatureMismatch)
}

pub fn derive_aes_key(license_id: &str) -> [u8; 32] {
    let hk = Hkdf::<Sha256>::new(None, EMBEDDED_AES_SALT);
    let mut okm = [0u8; 32];
    hk.expand(license_id.as_bytes(), &mut okm)
        .expect("HKDF expand for 32 bytes");
    okm
}

pub fn aes_gcm_decrypt(key: &[u8; 32], encrypted_data: &[u8]) -> Result<Vec<u8>, LicenseError> {
    if encrypted_data.len() < AES_NONCE_SIZE + 16 {
        return Err(LicenseError::DecryptionFailed);
    }

    let (nonce_bytes, ciphertext) = encrypted_data.split_at(AES_NONCE_SIZE);
    let nonce = Nonce::from_slice(nonce_bytes);
    let cipher =
        Aes256Gcm::new_from_slice(key).map_err(|_| LicenseError::CryptoError)?;
    cipher
        .decrypt(nonce, ciphertext)
        .map_err(|_| LicenseError::DecryptionFailed)
}

pub fn aes_gcm_encrypt(key: &[u8; 32], plaintext: &[u8]) -> Vec<u8> {
    use aes_gcm::aead::rand_core::RngCore;
    let mut nonce_bytes = [0u8; AES_NONCE_SIZE];
    OsRng.fill_bytes(&mut nonce_bytes);
    let nonce = Nonce::from_slice(&nonce_bytes);
    let cipher = Aes256Gcm::new_from_slice(key).expect("AES key valid");
    let mut result = nonce_bytes.to_vec();
    result.extend(cipher.encrypt(nonce, plaintext).expect("encrypt ok"));
    result
}

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

    #[test]
    fn test_derive_aes_key_deterministic() {
        let k1 = derive_aes_key("01JQABCDEFFGHIJKLMNOPQRSTUV");
        let k2 = derive_aes_key("01JQABCDEFFGHIJKLMNOPQRSTUV");
        assert_eq!(k1, k2);
    }

    #[test]
    fn test_derive_aes_key_different_ids() {
        let k1 = derive_aes_key("id-001");
        let k2 = derive_aes_key("id-002");
        assert_ne!(k1, k2);
    }

    #[test]
    fn test_aes_encrypt_decrypt_roundtrip() {
        let key = derive_aes_key("roundtrip-test");
        let plaintext = b"hello selfhosted licensing!";
        let encrypted = aes_gcm_encrypt(&key, plaintext);
        let decrypted = aes_gcm_decrypt(&key, &encrypted).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_aes_decrypt_wrong_key_fails() {
        let key1 = derive_aes_key("id-1");
        let key2 = derive_aes_key("id-2");
        let plaintext = b"secret data";
        let encrypted = aes_gcm_encrypt(&key1, plaintext);
        let result = aes_gcm_decrypt(&key2, &encrypted);
        assert!(result.is_err());
    }

    #[test]
    fn test_aes_decrypt_tampered_data_fails() {
        let key = derive_aes_key("tamper-test");
        let plaintext = b"original";
        let mut encrypted = aes_gcm_encrypt(&key, plaintext);
        if encrypted.len() > 20 {
            encrypted[20] ^= 0xFF;
        }
        let result = aes_gcm_decrypt(&key, &encrypted);
        assert!(result.is_err());
    }

    #[test]
    fn test_rsa_verify_with_generated_keypair() {
        use rsa::pkcs8::EncodePublicKey;
        use rsa::signature::Signer;
        use rsa::{pkcs8::LineEnding, RsaPrivateKey};
        use rand::rngs::OsRng;

        let mut rng = OsRng;
        let priv_key = RsaPrivateKey::new(&mut rng, 2048).unwrap();
        let pub_key = RsaPublicKey::from(&priv_key);
        let pub_pem = pub_key.to_public_key_pem(LineEnding::LF).unwrap();

        let data = b"test data to sign";
        let mut hasher = Sha256::new();
        hasher.update(data);
        let digest = hasher.finalize();
        let sig = priv_key
            .sign(Pkcs1v15Sign::new::<Sha256>(), &digest)
            .unwrap();
        let sig_bytes: &[u8] = sig.as_ref();
        let sig_b64 = URL_SAFE_NO_PAD.encode(sig_bytes);

        assert!(verify_rsa_signature(&pub_pem, data, &sig_b64).is_ok());
        assert!(verify_rsa_signature(&pub_pem, b"tampered", &sig_b64).is_err());
    }
}