#include <iostream>
#include <iomanip>
#include <openssl/aes.h>
#include <openssl/rand.h>
#include <cstring>
#include <chrono>
#include "aes.cpp"

int main() {
    const int numTests = 1000000;
    const int blockSize = 16;
    const int keySize = 16;
    uint8_t key[keySize];
    uint8_t opensslOutput[blockSize];
    uint8_t myOutput[blockSize];
    uint8_t input[blockSize];
    
    uint32_t expandedKey[44];

    std::chrono::duration<double, std::micro> opensslTime(0);
    std::chrono::duration<double, std::micro> myTime(0);

    AES_KEY opensslKey;

    for (int test = 0; test < numTests; ++test) {
        // Generate random key and input data using OpenSSL's random generator
        if (!RAND_bytes(key, keySize) || !RAND_bytes(input, blockSize)) {
            std::cerr << "Failed to generate random data using OpenSSL." << std::endl;
            return 1;
        }

        auto start = std::chrono::high_resolution_clock::now();
        AES_set_encrypt_key(key, 128, &opensslKey);
        AES_encrypt(input, opensslOutput, &opensslKey);
        auto end = std::chrono::high_resolution_clock::now();
        opensslTime += end - start;

        start = std::chrono::high_resolution_clock::now();
        expandKey(key, expandedKey);
        aes128(input, myOutput, expandedKey);
        end = std::chrono::high_resolution_clock::now();
        myTime += end - start;
        if (std::memcmp(opensslOutput, myOutput, blockSize))
            std::cerr << "Mismatch at test " << test + 1 << "!" << std::endl;


    }
    double avgOpenSSLTime = opensslTime.count() / numTests;
    double avgCustomTime = myTime.count() / numTests;
    
    std::cout << "avg openssl time: " << avgOpenSSLTime << "us" << std::endl;
    std::cout << "avg my time: " << avgCustomTime << "us" << std::endl;
    return 0;
}