#include <iostream>
#include <vector>
#include <chrono>
#include <iomanip>
#include "../src/crypto/aes.h"

void print_row(const std::string& name, size_t payload_size, double avg_time) {
    double throughput = (payload_size / (1024.0 * 1024.0)) / avg_time; // MB/s
    std::cout << std::left << std::setw(15) << name 
              << " | Size: " << std::setw(8) << payload_size << " B"
              << " | Avg: " << std::fixed << std::setprecision(6) << avg_time << " s"
              << " | Throughput: " << std::setprecision(2) << throughput << " MB/s" 
              << std::endl;
}

template<typename T>
void benchmark_aes_cbc(const std::string& name, const std::vector<uint8_t>& key, const std::vector<uint8_t>& iv, size_t payload_size) {
    T cipher(key);
    size_t padded_size = ((payload_size + 15) / 16) * 16;
    std::vector<uint8_t> pt(padded_size, 0xAA);
    
    auto ct = cipher.encryptCBC(pt, iv);

    const int iterations = 100;
    auto start = std::chrono::high_resolution_clock::now();
    for (int i = 0; i < iterations; ++i) {
        ct = cipher.encryptCBC(pt, iv);
    }
    auto end = std::chrono::high_resolution_clock::now();

    double avg_time = std::chrono::duration<double>(end - start).count() / iterations;
    print_row(name + "-CBC", payload_size, avg_time);
}

template<typename T>
void benchmark_aes_gcm(const std::string& name, const std::vector<uint8_t>& key, const uint8_t iv12[12], size_t payload_size) {
    T cipher(key);
    std::vector<uint8_t> pt(payload_size, 0xAA);
    std::vector<uint8_t> ct(payload_size);
    uint8_t tag[16];

    cipher.aes_gcm_encrypt(iv12, nullptr, 0, pt.data(), pt.size(), ct.data(), tag);

    const int iterations = 100;
    auto start = std::chrono::high_resolution_clock::now();
    for (int i = 0; i < iterations; ++i) {
        cipher.aes_gcm_encrypt(iv12, nullptr, 0, pt.data(), pt.size(), ct.data(), tag);
    }
    auto end = std::chrono::high_resolution_clock::now();

    double avg_time = std::chrono::duration<double>(end - start).count() / iterations;
    print_row(name + "-GCM", payload_size, avg_time);
}

int main() {
    std::vector<uint8_t> key128(16, 0x01);
    std::vector<uint8_t> key256(32, 0x01);
    std::vector<uint8_t> iv_cbc(16, 0x00);
    uint8_t iv_gcm[12] = {};

    std::cout << "AES Benchmarks\n";
    std::cout << "---------------------------------------------------------------------------\n";
    std::cout << std::left << std::setw(15) << "Algorithm" 
              << " | " << std::setw(16) << "Payload Size"
              << " | " << std::setw(15) << "Avg Time"
              << " | Throughput\n";
    std::cout << "---------------------------------------------------------------------------\n";

    std::vector<size_t> sizes = {64, 1024, 16384, 65536, 262144};

    for (auto size : sizes) {
        benchmark_aes_cbc<netplus::aes128>("AES-128", key128, iv_cbc, size);
        benchmark_aes_cbc<netplus::aes256>("AES-256", key256, iv_cbc, size);
        benchmark_aes_gcm<netplus::aes128>("AES-128", key128, iv_gcm, size);
        benchmark_aes_gcm<netplus::aes256>("AES-256", key256, iv_gcm, size);
    }

    std::cout << "---------------------------------------------------------------------------\n";

    return 0;
}
