use std::{fs::{File}, io::Write, path::Path};

use cas_lib::symmetric::{aes::CASAES256, cas_symmetric_encryption::CASAES256Encryption};

#[cfg(test)]
mod symmetric {
    use cas_lib::{key_exchange::{cas_key_exchange::CASKeyExchange, x25519::X25519}, pqc::{cas_pqc::MlKemKeyPair, ml_kem::{ml_kem_1024_decapsulate, ml_kem_1024_encapsulate, ml_kem_1024_generate}}, symmetric::{aes::CASAES128, cas_symmetric_encryption::{CASAES128Encryption, Chacha20Poly1305Encryption}, chacha20poly1305::CASChacha20Poly1305}};
    use hkdf::Hkdf;
    use sha2::Sha256;

    use super::*;
    #[test]
    fn test_aes_256() {
        let path = Path::new("tests/test.docx");
        let file_bytes: Vec<u8> = std::fs::read(path).unwrap();
        let aes_nonce = <CASAES256 as CASAES256Encryption>::generate_nonce();
        let aes_key = <CASAES256 as CASAES256Encryption>::generate_key();
        let encrypted_bytes = <CASAES256 as CASAES256Encryption>::encrypt_plaintext(aes_key.clone(), aes_nonce.clone(), file_bytes.clone());
        let mut file =  File::create("encrypted.docx").unwrap();
        file.write_all(&encrypted_bytes).unwrap();

        let decrypted_bytes = <CASAES256 as CASAES256Encryption>::decrypt_ciphertext(aes_key, aes_nonce, encrypted_bytes);
        let mut file =  File::create("decrypted.docx").unwrap();
        file.write_all(&decrypted_bytes).unwrap();
        assert_eq!(file_bytes, decrypted_bytes);
    }

    #[test]
    fn test_aes256_x25519_diffie_hellman() {
        let bob_public = <X25519 as CASKeyExchange>::generate_secret_and_public_key();
        let alice_public = <X25519 as CASKeyExchange>::generate_secret_and_public_key();

        let bob_shared_secret = <X25519 as CASKeyExchange>::diffie_hellman(bob_public.secret_key, alice_public.public_key);
        let alice_shared_secret = <X25519 as CASKeyExchange>::diffie_hellman(alice_public.secret_key, bob_public.public_key);

        let aes_nonce = <CASAES256 as CASAES256Encryption>::generate_nonce();
        let path = Path::new("tests/test.docx");
        let file_bytes: Vec<u8> = std::fs::read(path).unwrap();
        let alice_key = <CASAES256 as CASAES256Encryption>::key_from_x25519_shared_secret(bob_shared_secret);
        let encrypted_bytes = <CASAES256 as CASAES256Encryption>::encrypt_plaintext(alice_key, aes_nonce.clone(), file_bytes.clone());
        let mut file =  File::create("encrypted.docx").unwrap();
        file.write_all(&encrypted_bytes).unwrap();

        let bob_key = <CASAES256 as CASAES256Encryption>::key_from_x25519_shared_secret(alice_shared_secret);
        let decrypted_bytes = <CASAES256 as CASAES256Encryption>::decrypt_ciphertext(bob_key, aes_nonce, encrypted_bytes);
        let mut file =  File::create("decrypted.docx").unwrap();
        file.write_all(&decrypted_bytes).unwrap();
        assert_eq!(file_bytes, decrypted_bytes);
    }
    
    #[test]
    fn test_aes_128() {
        let path = Path::new("tests/test.docx");
        let file_bytes: Vec<u8> = std::fs::read(path).unwrap();
        let aes_nonce = <CASAES128 as CASAES128Encryption>::generate_nonce();
        let aes_key = <CASAES128 as CASAES128Encryption>::generate_key();
        let encrypted_bytes = <CASAES128 as CASAES128Encryption>::encrypt_plaintext(aes_key.clone(), aes_nonce.clone(), file_bytes.clone());
        let mut file =  File::create("encrypted.docx").unwrap();
        file.write_all(&encrypted_bytes).unwrap();


        let decrypted_bytes = <CASAES128 as CASAES128Encryption>::decrypt_ciphertext(aes_key, aes_nonce, encrypted_bytes);
        let mut file =  File::create("decrypted.docx").unwrap();
        file.write_all(&decrypted_bytes).unwrap();
        assert_eq!(file_bytes, decrypted_bytes);
    }

    #[test]
    fn test_aes128_x25519_diffie_hellman() {
        let bob_public = <X25519 as CASKeyExchange>::generate_secret_and_public_key();
        let alice_public = <X25519 as CASKeyExchange>::generate_secret_and_public_key();

        let bob_shared_secret = <X25519 as CASKeyExchange>::diffie_hellman(bob_public.secret_key, alice_public.public_key);
        let alice_shared_secret = <X25519 as CASKeyExchange>::diffie_hellman(alice_public.secret_key, bob_public.public_key);

        let aes_nonce = <CASAES128 as CASAES128Encryption>::generate_nonce();
        let path = Path::new("tests/test.docx");
        let file_bytes: Vec<u8> = std::fs::read(path).unwrap();
        let alice_key = <CASAES128 as CASAES128Encryption>::key_from_x25519_shared_secret(bob_shared_secret);
        let encrypted_bytes = <CASAES128 as CASAES128Encryption>::encrypt_plaintext(alice_key, aes_nonce.clone(), file_bytes.clone());
        let mut file =  File::create("encrypted.docx").unwrap();
        file.write_all(&encrypted_bytes).unwrap();
        

        let bob_key = <CASAES128 as CASAES128Encryption>::key_from_x25519_shared_secret(alice_shared_secret);
        let decrypted_bytes = <CASAES128 as CASAES128Encryption>::decrypt_ciphertext(bob_key, aes_nonce, encrypted_bytes);
        let mut file =  File::create("decrypted.docx").unwrap();
        file.write_all(&decrypted_bytes).unwrap();
        assert_eq!(file_bytes, decrypted_bytes);
    }

    #[test]
    fn test_chacha20_poly1305() {
        let path = Path::new("tests/test.docx");
        let file_bytes: Vec<u8> = std::fs::read(path).unwrap();
        let chacha20_nonce = <CASChacha20Poly1305 as Chacha20Poly1305Encryption>::generate_nonce();
        let chacha20_key = <CASChacha20Poly1305 as Chacha20Poly1305Encryption>::generate_key();
        let encrypted_bytes = <CASChacha20Poly1305 as Chacha20Poly1305Encryption>::encrypt_plaintext(chacha20_key.clone(), chacha20_nonce.clone(), file_bytes.clone());
        let mut file =  File::create("encrypted.docx").unwrap();
        file.write_all(&encrypted_bytes).unwrap();

        let decrypted_bytes = <CASChacha20Poly1305 as Chacha20Poly1305Encryption>::decrypt_ciphertext(chacha20_key, chacha20_nonce, encrypted_bytes);
        let mut file =  File::create("decrypted.docx").unwrap();
        file.write_all(&decrypted_bytes).unwrap();
        assert_eq!(file_bytes, decrypted_bytes);
    }

    #[test]
    pub fn round_trip_mlkem1024_hkdf() {
        let secret_key_public_key: MlKemKeyPair = ml_kem_1024_generate();
        let ct = ml_kem_1024_encapsulate(secret_key_public_key.public_key).expect("encapsulate failed");
        let ss_receiver = ml_kem_1024_decapsulate(secret_key_public_key.secret_key, ct.ciphertext).expect("decapsulate failed");

        let bob_shared_secret = Hkdf::<Sha256>::new(None, &ss_receiver);
        let alice_shared_secret = Hkdf::<Sha256>::new(None, &ct.shared_secret);

        let mut aes_key = Box::new([0u8; 32]);
        bob_shared_secret.expand(b"aes key", &mut *aes_key).unwrap();

        let mut aes_key2 = Box::new([0u8; 32]);
        alice_shared_secret.expand(b"aes key", &mut *aes_key2).unwrap();
        assert_eq!(aes_key.to_vec(), aes_key2.to_vec());

        let aes_nonce = <CASAES256 as CASAES256Encryption>::generate_nonce();
        let path = Path::new("tests/test.docx");
        let file_bytes: Vec<u8> = std::fs::read(path).unwrap();
        let encrypted_bytes = <CASAES256 as CASAES256Encryption>::encrypt_plaintext(aes_key.to_vec(), aes_nonce.clone(), file_bytes.clone());
        let mut file =  File::create("encrypted.docx").unwrap();
        file.write_all(&encrypted_bytes).unwrap();

        let decrypted_bytes = <CASAES256 as CASAES256Encryption>::decrypt_ciphertext(aes_key2.to_vec(), aes_nonce, encrypted_bytes);
        let mut file =  File::create("decrypted.docx").unwrap();
        file.write_all(&decrypted_bytes).unwrap();
        assert_eq!(file_bytes, decrypted_bytes);
    }
}