use rsa::{RsaPrivateKey, RsaPublicKey, Pkcs1v15Encrypt};
use std::time::Instant;
use rand::rngs::OsRng;
use rand::Rng;
use std::fs::File;
use std::error::Error;
use aes::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use des::TdesEde3;
use cbc::{Encryptor, Decryptor};

type DesEde3CbcEnc = Encryptor<TdesEde3>;
type DesEde3CbcDec = Decryptor<TdesEde3>;

type Aes128CbcEnc = cbc::Encryptor<aes::Aes128>;
type Aes128CbcDec = cbc::Decryptor<aes::Aes128>;
type Aes192CbcEnc = cbc::Encryptor<aes::Aes192>;
type Aes192CbcDec = cbc::Decryptor<aes::Aes192>;
type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;
type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;


fn generate_rsa_keys(count: usize, key_size: usize, mut rng: OsRng) -> Vec<(RsaPrivateKey, RsaPublicKey)> {
    (0..count)
        .map(|_| {
            let private_key = RsaPrivateKey::new(&mut rng, key_size).expect("Failed to generate a key");
            let public_key = RsaPublicKey::from(&private_key);
            (private_key, public_key)
        })
        .collect()
}

fn generate_3des_keys(key_size_bits: usize, num_keys: usize, mut rng: OsRng) -> Vec<Vec<u8>> {
    let key_size_bytes = key_size_bits / 8;

    (0..num_keys)
        .map(|_| {
            let mut key = vec![0u8; key_size_bytes];
            rng.fill(&mut key[..]);
            key
        })
        .collect()
}

fn generate_aes_keys(key_size_bits: usize, num_keys: usize, mut rng: OsRng) -> Vec<Vec<u8>> {
    let key_size_bytes = key_size_bits / 8;

    (0..num_keys)
        .map(|_| {
            let mut key = vec![0u8; key_size_bytes];
            rng.fill(&mut key[..]);
            key
        })
        .collect()
}

fn split_bytes_into_chunks(input: &[u8], key_length: usize) -> Vec<&[u8]> {
    // RSA key length determines the maximum chunk size we can encrypt.
    // Adjust for padding overhead (e.g., for PKCS#1v15, typically 11 bytes).
    let max_chunk_size = key_length/8 - 11; // Example: PKCS#1 padding overhead

    // Use `chunks` to split the slice into smaller chunks and collect into a Vec of slices
    input.chunks(max_chunk_size).collect()
}

fn encrypt_3des(key: &[u8], iv: &[u8], plaintext: &[u8]) -> Vec<u8> {
    let encryptor = DesEde3CbcEnc::new_from_slices(key, iv).expect("Preparing encryptor error");
    let ciphertext = encryptor.encrypt_padded_vec_mut::<Pkcs7>(plaintext);
    ciphertext
}

fn decrypt_3des(key: &[u8], iv: &[u8], ciphertext: &[u8]) -> Vec<u8> {
    let decryptor = DesEde3CbcDec::new_from_slices(key, iv).expect("Preparing decryptor error");
    let decrypted = decryptor.decrypt_padded_vec_mut::<Pkcs7>(ciphertext).expect("Error decrypting");
    decrypted
}

fn aes_encrypt(key: &[u8], iv: &[u8], plaintext: &[u8]) -> Vec<u8> {
    match key.len() {
        16 => Aes128CbcEnc::new_from_slices(key, iv)
            .unwrap()
            .encrypt_padded_vec_mut::<Pkcs7>(plaintext),
        24 => Aes192CbcEnc::new_from_slices(key, iv)
            .unwrap()
            .encrypt_padded_vec_mut::<Pkcs7>(plaintext),
        32 => Aes256CbcEnc::new_from_slices(key, iv)
            .unwrap()
            .encrypt_padded_vec_mut::<Pkcs7>(plaintext),
        _ => panic!("Invalid key length. Must be 16, 24, or 32 bytes."),
    }
}

fn aes_decrypt(key: &[u8], iv: &[u8], ciphertext: &[u8]) -> Vec<u8> {
    match key.len() {
        16 => Aes128CbcDec::new_from_slices(key, iv)
            .unwrap()
            .decrypt_padded_vec_mut::<Pkcs7>(ciphertext)
            .unwrap(),
        24 => Aes192CbcDec::new_from_slices(key, iv)
            .unwrap()
            .decrypt_padded_vec_mut::<Pkcs7>(ciphertext)
            .unwrap(),
        32 => Aes256CbcDec::new_from_slices(key, iv)
            .unwrap()
            .decrypt_padded_vec_mut::<Pkcs7>(ciphertext)
            .unwrap(),
        _ => panic!("Invalid key length. Must be 16, 24, or 32 bytes."),
    }
}

fn measure_rsa_performance(
    key: &RsaPrivateKey,
    public_key: &RsaPublicKey,
    key_size: usize,
    sizes: &[usize]
) -> Result<(), Box<dyn Error>> {
    println!("Measuting rsa{} performance", key_size.to_string());
    let mut encryption_results = vec![];
    let mut decryption_results = vec![];

    for &size in sizes {
        let plaintext: Vec<u8> = vec![0u8; size];
        
        println!("First 10 bytes of plaintext: {:02X?}", &plaintext[..10]);
        
        let data_vec: Vec<&[u8]> = split_bytes_into_chunks(&plaintext, key_size);
        println!("Splitted into chunks: {}", data_vec.len());
        let mut ciphertexts: Vec<Vec<u8>> = Vec::new();
        let mut decrypted_texts: Vec<Vec<u8>> = Vec::new();

        // Measure encryption time
        let start = Instant::now();
        for data in data_vec {
            let ciphertext = public_key.encrypt(
                &mut OsRng,
                Pkcs1v15Encrypt,
                &data
            )?;
            ciphertexts.push(ciphertext);
        }
        let encrypt_duration = start.elapsed();
        println!("First 10 bytes of ciphertext: {:02X?}", &ciphertexts[0][..10]);

        // Measure decryption time
        let start = Instant::now();
        for data in ciphertexts {
            let decrypted = key.decrypt(
                Pkcs1v15Encrypt,
                &data
            )?;
            decrypted_texts.push(decrypted);
        }
        let decrypt_duration = start.elapsed();
        println!("First 10 bytes of decrypted plaintext: {:02X?}", &decrypted_texts[0][..10]);

        // Record results
        encryption_results.push((size, encrypt_duration.as_micros() as f32 / 1000.0));
        decryption_results.push((size, decrypt_duration.as_micros() as f32 / 1000.0));
    }

    // Save encryption results to CSV
    let encrypt_file = File::create(format!("rsa{}_encryption.csv", key_size))?;
    let mut wtr = csv::Writer::from_writer(encrypt_file);
    wtr.write_record(&["key_size", "data_size", "time_ms"])?;
    for (size, time) in encryption_results {
        wtr.write_record(&[key_size.to_string(), size.to_string(), time.to_string()])?;
    }

    // Save decryption results to CSV
    let decrypt_file = File::create(format!("rsa{}_decryption.csv", key_size))?;
    let mut wtr = csv::Writer::from_writer(decrypt_file);
    wtr.write_record(&["key_size", "data_size", "time_ms"])?;
    for (size, time) in decryption_results {
        wtr.write_record(&[key_size.to_string(), size.to_string(), time.to_string()])?;
    }

    Ok(())
}

fn measure_aes_performance(
        key: &[u8],
        iv: &[u8],
        sizes: &[usize]
    ) -> Result<(), Box<dyn Error>> {
    println!("Measuting aes{} performance", (key.len()*8).to_string());
    let mut encryption_results = vec![];
    let mut decryption_results = vec![];

    for &size in sizes {
        let plaintext: Vec<u8> = vec![0u8; size];

        println!("First 10 bytes of plaintext: {:02X?}", &plaintext[..10]);
        // Measure encryption time
        let start = Instant::now();
        let ciphertext = aes_encrypt(key, iv, &plaintext);
        let encrypt_duration = start.elapsed();
        println!("First 10 bytes of ciphertext: {:02X?}", &ciphertext[..10]);

        // Measure decryption time
        let start = Instant::now();
        let decrypted = aes_decrypt(key, iv, &ciphertext);
        let decrypt_duration = start.elapsed();
        println!("First 10 bytes of decrypted plaintext: {:02X?}", &decrypted[..10]);

        // Record results
        encryption_results.push((size, encrypt_duration.as_micros() as f32 / 1000.0));
        decryption_results.push((size, decrypt_duration.as_micros() as f32 / 1000.0));
    }

    let encrypt_file = File::create(format!("aes{}_encryption.csv", key.len()*8))?;
    let mut wtr = csv::Writer::from_writer(encrypt_file);
    wtr.write_record(&["key_size", "data_size", "time_ms"])?;
    for (size, time) in encryption_results {
        wtr.write_record(&[(key.len()*8).to_string(), size.to_string(), time.to_string()])?;
    }

    // Save decryption results to CSV
    let decrypt_file = File::create(format!("aes{}_decryption.csv", key.len()*8))?;
    let mut wtr = csv::Writer::from_writer(decrypt_file);
    wtr.write_record(&["key_size", "data_size", "time_ms"])?;
    for (size, time) in decryption_results {
        wtr.write_record(&[(key.len()*8).to_string(), size.to_string(), time.to_string()])?;
    }
    Ok(())
}

fn measure_3des_performance(
    key: &[u8],
    iv: &[u8],
    sizes: &[usize]
) -> Result<(), Box<dyn Error>> {
    let mut encryption_results = vec![];
    let mut decryption_results = vec![];
    println!("Measuring 3des performance...");

    for &size in sizes {
        let plaintext: Vec<u8> = vec![0u8; size];

        println!("First 10 bytes of plaintext: {:02X?}", &plaintext[..10]);

        // Measure encryption time
        let start = Instant::now();
        let ciphertext = encrypt_3des(&key, &iv, &plaintext);
        let encrypt_duration = start.elapsed();
        println!("First 10 bytes of ciphertext: {:02X?}", &ciphertext[..10]);

        // Measure decryption time
        let start = Instant::now();
        let decrypted = decrypt_3des(&key, &iv, &ciphertext);
        let decrypt_duration = start.elapsed();
        println!("First 10 bytes of decrypted plaintext: {:02X?}", &decrypted[..10]);

        // Record results
        encryption_results.push((size, encrypt_duration.as_micros() as f32 / 1000.0));
        decryption_results.push((size, decrypt_duration.as_micros() as f32 / 1000.0));
    }

    // Save encryption results to CSV
    let encrypt_file = File::create(format!("3des{}_encryption.csv", key.len() * 8))?;
    let mut wtr = csv::Writer::from_writer(encrypt_file);
    wtr.write_record(&["key_size", "data_size", "time_ms"])?;
    for (size, time) in encryption_results {
        wtr.write_record(&[(key.len()*8).to_string(), size.to_string(), time.to_string()])?;
    }

    // Save decryption results to CSV
    let decrypt_file = File::create(format!("3des{}_decryption.csv", key.len() * 8))?;
    let mut wtr = csv::Writer::from_writer(decrypt_file);
    wtr.write_record(&["key_size", "data_size", "time_ms"])?;
    for (size, time) in decryption_results {
        wtr.write_record(&[(key.len()*8).to_string(), size.to_string(), time.to_string()])?;
    }

    Ok(())
}

fn test_rsa_key_gen(rng: OsRng) -> Result<(), Box<dyn Error>> {
    let counts = [1, 10, 100, 1000 ];
    let key_sizes = [512, 1024, 2048 ];

    let file = File::create("rsa_keygen_benchmark.csv")?;
    let mut wtr = csv::Writer::from_writer(file);

    wtr.write_record(&["count", "key_size", "time_ms"])?;

    for &count in &counts {
        for &key_size in &key_sizes {

            let start = Instant::now();
            generate_rsa_keys(count, key_size, rng);
            let duration = start.elapsed();

            let time_us = duration.as_micros();

            wtr.write_record(&[
                count.to_string(),
                key_size.to_string(),
                ((time_us as f32)/1000.0).to_string(),
            ])?;
        }
    }
    wtr.flush()?;
    Ok(())
}

fn test_3des_key_gen(rng: OsRng) -> Result<(), Box<dyn Error>> {
    let counts = [1, 10, 100, 1000 ];
    let key_sizes = [192 ];

    let file = File::create("3des_keygen_benchmark.csv")?;
    let mut wtr = csv::Writer::from_writer(file);

    wtr.write_record(&["count", "key_size", "time_ms"])?;

    for &count in &counts {
        for &key_size in &key_sizes {

            let start = Instant::now();
            generate_3des_keys(count, key_size, rng);
            let duration = start.elapsed();

            let time_us = duration.as_micros();

            wtr.write_record(&[
                count.to_string(),
                key_size.to_string(),
                ((time_us as f32)/1000.0).to_string(),
            ])?;
        }
    }
    wtr.flush()?;
    Ok(())
}

fn test_aes_key_gen(rng: OsRng) -> Result<(), Box<dyn Error>> {
    let counts = [1, 10, 100, 1000 ];
    let key_sizes = [128, 192, 256 ];

    let file = File::create("aes_keygen_benchmark.csv")?;
    let mut wtr = csv::Writer::from_writer(file);

    wtr.write_record(&["count", "key_size", "time_ms"])?;

    for &count in &counts {
        for &key_size in &key_sizes {

            let start = Instant::now();
            generate_aes_keys(count, key_size, rng);
            let duration = start.elapsed();

            let time_us = duration.as_micros();

            wtr.write_record(&[
                count.to_string(),
                key_size.to_string(),
                ((time_us as f32)/1000.0).to_string(),
            ])?;
        }
    }
    wtr.flush()?;
    Ok(())
}

fn main() {
    let rng = OsRng;
    let _result = test_rsa_key_gen(rng);
    let _result = test_3des_key_gen(rng);
    let _result = test_aes_key_gen(rng);

    let iv = [0x8; 16];
    let data_sizes = [10, 1000, 100000, 10000000 ];

    let _result = measure_aes_performance(&generate_aes_keys(128, 1, rng)[0], &iv, &data_sizes);
    let _result = measure_aes_performance(&generate_aes_keys(192, 1, rng)[0], &iv, &data_sizes);
    let _result = measure_aes_performance(&generate_aes_keys(256, 1, rng)[0], &iv, &data_sizes);

    let iv = [0x8; 8];
    let _result = measure_3des_performance(&generate_3des_keys(192, 1, rng)[0], &iv, &data_sizes);

    for size in [512, 1024, 2048, 3072] {
        let keys = generate_rsa_keys(1, size, rng);
        let (privkey, pubkey) = &keys[0];
        match measure_rsa_performance(&privkey, &pubkey, size, &data_sizes) {
            Ok(_) => println!("Function executed successfully!"),
            Err(e) => println!("Error occurred: {}", e),
        }
    }
}