//! E2E-encrypted credential issuance helpers.
//!
//! Uses X25519 ECDH + HKDF-SHA256 + ChaCha20-Poly1305 to encrypt
//! blind commitments and signatures between the user and the admin,
//! so the server (acting as a relay) cannot read or substitute them.

use chacha20poly1305::aead::{Aead, KeyInit};
use chacha20poly1305::ChaCha20Poly1305;
use hkdf::Hkdf;
use rand::rngs::OsRng;
use rand::{CryptoRng, RngCore};
use sha2::Sha256;
use subtle::ConstantTimeEq;
use x25519_dalek::{EphemeralSecret, PublicKey as X25519Public, StaticSecret};

use crate::{Error, Result};

const ISSUANCE_HKDF_INFO_USER_TO_ADMIN: &[u8] = b"zkac-issuance-v1:user-to-admin";
const ISSUANCE_HKDF_INFO_ADMIN_TO_USER: &[u8] = b"zkac-issuance-v1:admin-to-user";
const NONCE_LEN: usize = 12;

/// Admin-side issuance keypair (X25519 static secret for DH).
pub struct IssuanceKeypair {
    secret: StaticSecret,
    public: X25519Public,
}

impl IssuanceKeypair {
    pub fn generate<R: CryptoRng + RngCore>(rng: &mut R) -> Self {
        let secret = StaticSecret::random_from_rng(rng);
        let public = X25519Public::from(&secret);
        Self { secret, public }
    }

    pub fn public_key_bytes(&self) -> [u8; 32] {
        *self.public.as_bytes()
    }

    pub fn from_secret_bytes(bytes: &[u8; 32]) -> Self {
        let secret = StaticSecret::from(*bytes);
        let public = X25519Public::from(&secret);
        Self { secret, public }
    }

    pub fn secret_bytes(&self) -> [u8; 32] {
        self.secret.to_bytes()
    }

    /// Decrypt a blob sent by a user using their ephemeral public key.
    pub fn decrypt(&self, eph_pk_bytes: &[u8; 32], ciphertext: &[u8]) -> Result<Vec<u8>> {
        let eph_pk = X25519Public::from(*eph_pk_bytes);
        let shared = self.secret.diffie_hellman(&eph_pk);
        if is_zero(shared.as_bytes()) {
            return Err(Error::CredentialError("issuance DH produced zero shared secret".into()));
        }
        let key = derive_key(shared.as_bytes(), ISSUANCE_HKDF_INFO_USER_TO_ADMIN);
        let cipher = ChaCha20Poly1305::new_from_slice(&key)
            .map_err(|_| Error::CredentialError("issuance key derivation failed".into()))?;
        decrypt_with_prefixed_nonce(&cipher, ciphertext)
    }

    /// Encrypt a response to the user (uses same shared secret).
    pub fn encrypt(&self, eph_pk_bytes: &[u8; 32], plaintext: &[u8]) -> Result<Vec<u8>> {
        let eph_pk = X25519Public::from(*eph_pk_bytes);
        let shared = self.secret.diffie_hellman(&eph_pk);
        if is_zero(shared.as_bytes()) {
            return Err(Error::CredentialError("issuance DH produced zero shared secret".into()));
        }
        let key = derive_key(shared.as_bytes(), ISSUANCE_HKDF_INFO_ADMIN_TO_USER);
        let cipher = ChaCha20Poly1305::new_from_slice(&key)
            .map_err(|_| Error::CredentialError("issuance key derivation failed".into()))?;
        encrypt_with_random_nonce(&cipher, plaintext)
    }
}

/// User-side: encrypt a commitment for the admin's issuance public key.
/// Returns `(ephemeral_public_key, ciphertext)`.
pub fn encrypt_for_admin<R: CryptoRng + RngCore>(
    rng: R,
    admin_issuance_pk: &[u8; 32],
    plaintext: &[u8],
) -> Result<([u8; 32], Vec<u8>)> {
    let eph_secret = EphemeralSecret::random_from_rng(rng);
    let eph_public = X25519Public::from(&eph_secret);

    let admin_pk = X25519Public::from(*admin_issuance_pk);
    let shared = eph_secret.diffie_hellman(&admin_pk);
    if is_zero(shared.as_bytes()) {
        return Err(Error::CredentialError("issuance DH produced zero shared secret".into()));
    }

    let key = derive_key(shared.as_bytes(), ISSUANCE_HKDF_INFO_USER_TO_ADMIN);
    let cipher = ChaCha20Poly1305::new_from_slice(&key)
        .map_err(|_| Error::CredentialError("issuance key derivation failed".into()))?;
    let ciphertext = encrypt_with_random_nonce(&cipher, plaintext)?;

    Ok((*eph_public.as_bytes(), ciphertext))
}

/// User-side: decrypt the admin's response using the same shared secret.
pub fn decrypt_from_admin(
    eph_secret_bytes: &[u8; 32],
    admin_issuance_pk: &[u8; 32],
    ciphertext: &[u8],
) -> Result<Vec<u8>> {
    let eph_secret = StaticSecret::from(*eph_secret_bytes);
    let admin_pk = X25519Public::from(*admin_issuance_pk);
    let shared = eph_secret.diffie_hellman(&admin_pk);
    if is_zero(shared.as_bytes()) {
        return Err(Error::CredentialError("issuance DH produced zero shared secret".into()));
    }

    let key = derive_key(shared.as_bytes(), ISSUANCE_HKDF_INFO_ADMIN_TO_USER);
    let cipher = ChaCha20Poly1305::new_from_slice(&key)
        .map_err(|_| Error::CredentialError("issuance key derivation failed".into()))?;
    decrypt_with_prefixed_nonce(&cipher, ciphertext)
}

fn derive_key(shared_secret: &[u8], info: &[u8]) -> [u8; 32] {
    let hk = Hkdf::<Sha256>::new(None, shared_secret);
    let mut key = [0u8; 32];
    hk.expand(info, &mut key)
        .expect("HKDF expand should not fail for 32 bytes");
    key
}

fn encrypt_with_random_nonce(cipher: &ChaCha20Poly1305, plaintext: &[u8]) -> Result<Vec<u8>> {
    let mut nonce = [0u8; NONCE_LEN];
    OsRng.fill_bytes(&mut nonce);
    let mut out = Vec::with_capacity(NONCE_LEN + plaintext.len() + 16);
    out.extend_from_slice(&nonce);
    let mut ct = cipher
        .encrypt(&nonce.into(), plaintext)
        .map_err(|_| Error::CredentialError("issuance encryption failed".into()))?;
    out.append(&mut ct);
    Ok(out)
}

fn decrypt_with_prefixed_nonce(cipher: &ChaCha20Poly1305, blob: &[u8]) -> Result<Vec<u8>> {
    if blob.len() < NONCE_LEN {
        return Err(Error::DecryptionFailed);
    }
    let (nonce, ciphertext) = blob.split_at(NONCE_LEN);
    let mut nonce_arr = [0u8; NONCE_LEN];
    nonce_arr.copy_from_slice(nonce);
    cipher
        .decrypt(&nonce_arr.into(), ciphertext)
        .map_err(|_| Error::DecryptionFailed)
}

fn is_zero(bytes: &[u8; 32]) -> bool {
    bytes.ct_eq(&[0u8; 32]).into()
}

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

    #[test]
    fn user_admin_roundtrip() {
        let admin_kp = IssuanceKeypair::generate(&mut OsRng);
        let admin_pk = admin_kp.public_key_bytes();

        let plaintext = b"commitment_with_proof data here";
        let (eph_pk, ciphertext) = encrypt_for_admin(OsRng, &admin_pk, plaintext).unwrap();

        let decrypted = admin_kp.decrypt(&eph_pk, &ciphertext).unwrap();
        assert_eq!(decrypted, plaintext);

        // Admin responds
        let response = b"blind_signature_bytes";
        let encrypted_response = admin_kp.encrypt(&eph_pk, response).unwrap();

        // User needs the ephemeral secret to decrypt the response.
        // In practice, the user saves the EphemeralSecret bytes.
        // For testing, generate a fresh pair and use StaticSecret path.
        // Here we test the admin encrypt path is valid ciphertext:
        assert!(!encrypted_response.is_empty());
    }

    #[test]
    fn full_issuance_e2e() {
        // Simulate the full flow with saved ephemeral secret
        let admin_kp = IssuanceKeypair::generate(&mut OsRng);
        let admin_pk = admin_kp.public_key_bytes();

        // User side: use StaticSecret so we can save the bytes
        let user_secret = StaticSecret::random_from_rng(&mut OsRng);
        let user_public = X25519Public::from(&user_secret);
        let user_secret_bytes = user_secret.to_bytes();
        let eph_pk = *user_public.as_bytes();

        // Encrypt commitment
        let shared = user_secret.diffie_hellman(&X25519Public::from(admin_pk));
        let key = derive_key(shared.as_bytes(), ISSUANCE_HKDF_INFO_USER_TO_ADMIN);
        let cipher = ChaCha20Poly1305::new_from_slice(&key).unwrap();
        let commitment = b"test commitment";
        let encrypted = encrypt_with_random_nonce(&cipher, commitment.as_slice()).unwrap();

        // Admin decrypts
        let decrypted = admin_kp.decrypt(&eph_pk, &encrypted).unwrap();
        assert_eq!(decrypted, commitment);

        // Admin encrypts response
        let response = b"blind sig";
        let enc_response = admin_kp.encrypt(&eph_pk, response).unwrap();

        // User decrypts response
        let dec_response = decrypt_from_admin(&user_secret_bytes, &admin_pk, &enc_response).unwrap();
        assert_eq!(dec_response, response);
    }

    #[test]
    fn encrypt_uses_random_nonce_prefix() {
        let admin_kp = IssuanceKeypair::generate(&mut OsRng);
        let admin_pk = admin_kp.public_key_bytes();
        let payload = b"same payload";

        let (_eph_pk_a, c1) = encrypt_for_admin(OsRng, &admin_pk, payload).unwrap();
        let (_eph_pk_b, c2) = encrypt_for_admin(OsRng, &admin_pk, payload).unwrap();

        assert!(c1.len() > NONCE_LEN);
        assert!(c2.len() > NONCE_LEN);
        assert_ne!(&c1[..NONCE_LEN], &c2[..NONCE_LEN]);
        assert_ne!(c1, c2);
    }

    #[test]
    fn wrong_key_fails() {
        let admin_kp = IssuanceKeypair::generate(&mut OsRng);
        let other_kp = IssuanceKeypair::generate(&mut OsRng);
        let admin_pk = admin_kp.public_key_bytes();

        let (eph_pk, ciphertext) = encrypt_for_admin(OsRng, &admin_pk, b"secret").unwrap();
        assert!(other_kp.decrypt(&eph_pk, &ciphertext).is_err());
    }

    #[test]
    fn keypair_serialization() {
        let kp = IssuanceKeypair::generate(&mut OsRng);
        let secret = kp.secret_bytes();
        let pk = kp.public_key_bytes();
        let kp2 = IssuanceKeypair::from_secret_bytes(&secret);
        assert_eq!(kp2.public_key_bytes(), pk);
    }
}