//! OSTP Crypto Module - PSK, X25519, ChaCha20-Poly1305

use chacha20poly1305::{
    aead::{Aead, KeyInit, OsRng},
    ChaCha20Poly1305, Nonce,
};
use hmac::{digest::KeyInit as HmacKeyInit, Hmac, Mac};
use sha2::Sha256;
use x25519_dalek::{EphemeralSecret, PublicKey, SharedSecret};

/// AEAD error type re-export
pub type AeadError = chacha20poly1305::aead::Error;

type HmacSha256 = Hmac<Sha256>;

/// PSK-based silent handshake validator
pub struct PskValidator {
    psk: [u8; 32],
}

impl PskValidator {
    pub fn new(psk: [u8; 32]) -> Self {
        Self { psk }
    }

    /// Generate HMAC signature for packet authentication
    pub fn sign(&self, data: &[u8]) -> [u8; 32] {
        let mut mac: HmacSha256 = HmacKeyInit::new_from_slice(&self.psk).expect("HMAC key size");
        mac.update(data);
        mac.finalize().into_bytes().into()
    }

    /// Verify packet has valid PSK-derived signature (silent drop if invalid)
    pub fn verify(&self, data: &[u8], signature: &[u8; 32]) -> bool {
        let mut mac: HmacSha256 = HmacKeyInit::new_from_slice(&self.psk).expect("HMAC key size");
        mac.update(data);
        mac.verify_slice(signature).is_ok()
    }
}

/// X25519 key exchange
pub struct KeyExchange {
    secret: EphemeralSecret,
    public: PublicKey,
}

impl KeyExchange {
    pub fn new() -> Self {
        let secret = EphemeralSecret::random_from_rng(OsRng);
        let public = PublicKey::from(&secret);
        Self { secret, public }
    }

    pub fn public_key(&self) -> &PublicKey {
        &self.public
    }

    pub fn derive_shared(self, peer_public: &PublicKey) -> SharedSecret {
        self.secret.diffie_hellman(peer_public)
    }
}

impl Default for KeyExchange {
    fn default() -> Self {
        Self::new()
    }
}

/// AEAD cipher for data encryption
pub struct AeadCipher {
    cipher: ChaCha20Poly1305,
    nonce_counter: u64,
}

impl AeadCipher {
    pub fn new(key: &[u8; 32]) -> Self {
        let cipher = ChaCha20Poly1305::new_from_slice(key).expect("key size");
        Self {
            cipher,
            nonce_counter: 0,
        }
    }

    fn next_nonce(&mut self) -> Nonce {
        let mut nonce = [0u8; 12];
        nonce[4..12].copy_from_slice(&self.nonce_counter.to_le_bytes());
        self.nonce_counter = self.nonce_counter.wrapping_add(1);
        Nonce::from(nonce)
    }

    pub fn encrypt(&mut self, plaintext: &[u8]) -> Result<Vec<u8>, AeadError> {
        let nonce = self.next_nonce();
        let mut ciphertext = self.cipher.encrypt(&nonce, plaintext)?;
        // Prepend nonce for decryption
        let mut output = nonce.to_vec();
        output.append(&mut ciphertext);
        Ok(output)
    }

    pub fn decrypt(&self, ciphertext: &[u8]) -> Result<Vec<u8>, AeadError> {
        if ciphertext.len() < 12 {
            return Err(chacha20poly1305::aead::Error);
        }
        let nonce = Nonce::from_slice(&ciphertext[..12]);
        self.cipher.decrypt(nonce, &ciphertext[12..])
    }
}

/// Derive session key from shared secret using HKDF-like expansion
pub fn derive_session_key(shared: &SharedSecret, salt: &[u8]) -> [u8; 32] {
    let mut mac: HmacSha256 = HmacKeyInit::new_from_slice(salt).expect("HMAC key size");
    mac.update(shared.as_bytes());
    mac.finalize().into_bytes().into()
}

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

    #[test]
    fn test_psk_sign_verify() {
        let psk = [0x42u8; 32];
        let validator = PskValidator::new(psk);
        let data = b"hello";
        let sig = validator.sign(data);
        assert!(validator.verify(data, &sig));
    }

    #[test]
    fn test_key_exchange() {
        let alice = KeyExchange::new();
        let bob = KeyExchange::new();
        let alice_pub = *alice.public_key();
        let bob_pub = *bob.public_key();
        let alice_shared = alice.derive_shared(&bob_pub);
        let bob_shared = bob.derive_shared(&alice_pub);
        assert_eq!(alice_shared.as_bytes(), bob_shared.as_bytes());
    }

    #[test]
    fn test_aead_encrypt_decrypt() {
        let key = [0x11u8; 32];
        let mut cipher = AeadCipher::new(&key);
        let plaintext = b"secret message";
        let ciphertext = cipher.encrypt(plaintext).unwrap();
        let decrypted = cipher.decrypt(&ciphertext).unwrap();
        assert_eq!(plaintext.as_slice(), decrypted.as_slice());
    }
}
