#include <aes.hpp>
#include <fstream>
#include <string>
#include <util.h>

#ifdef USE_BOOST
#include <boost/filesystem.hpp>
namespace fs = boost::filesystem;
#else
#include <filesystem>
namespace fs = std::filesystem;
#endif

const size_t BLOCK_SIZE = 102400;   // 100KB块大小
const size_t IV_SIZE = 16;          // 随机值大小

/*
测试环境：
Microsoft Windows 10 Professional (x64) Build 19045.5608 (22H2)
13th Gen Intel(R) Core(TM) i5-13500H 3200.0 MHz
Debug模式 tinyaes 加密解密测试速度：
=========================================
File size: 630239232 bytes (601.043 MB)
Effective block size: 102384 bytes
IV size: 16 bytes
Blocks processed: 6156
Total encryption time: 41887336 μs (14.349 MB/s)
Total decryption time: 41822620 μs (14.3712 MB/s)
Decrypted file: "D:\\1_decrypted.iso"
Data verification: PASSED
=========================================
Release模式 tinyaes 加密解密测试速度：
=========================================
File size: 630239232 bytes (601.043 MB)
Effective block size: 102384 bytes
IV size: 16 bytes
Blocks processed: 6156
Total encryption time: 8367460 μs (71.831 MB/s)
Total decryption time: 8150036 μs (73.7473 MB/s)
Decrypted file: "D:\\1_decrypted.iso"
Data verification: PASSED
=========================================
*/

void test_speed(const char* input_file)
{
    // 检查输入文件
    if (!fs::exists(input_file)) {
        std::cerr << "Input file not found: " << input_file << std::endl;
        return;
    }

    size_t file_size = fs::file_size(input_file);
    if (file_size == 0) {
        std::cerr << "Input file is empty" << std::endl;
        return;
    }

    // 准备密钥
    std::string key = "test_speed_key";
    uint8_t ik[32]{};
    hash(key.c_str(), ik);

    // 准备解密后的输出文件
    fs::path decrypted_path = fs::path(input_file)
                                  .replace_filename(fs::path(input_file).stem().string() + "_decrypted" +
                                                    fs::path(input_file).extension().string());

    std::ofstream decrypted_file(decrypted_path, std::ios::binary);
    if (!decrypted_file) {
        std::cerr << "Failed to create decrypted file" << std::endl;
        return;
    }

    // 打开输入文件
    std::ifstream in_file(input_file, std::ios::binary);
    if (!in_file) {
        std::cerr << "Failed to open input file" << std::endl;
        return;
    }

    // 测试数据缓冲区（额外预留16字节空间）
    std::vector<uint8_t> original_block(BLOCK_SIZE);
    std::vector<uint8_t> processing_block(BLOCK_SIZE + IV_SIZE);   // 加密/解密处理缓冲区

    size_t total_bytes = 0;
    size_t blocks_processed = 0;
    bool verification_passed = true;

    // 计时变量
    auto total_encrypt_time = std::chrono::microseconds(0);
    auto total_decrypt_time = std::chrono::microseconds(0);

    while (in_file) {
        // 读取原始数据块（注意实际读取量比BLOCK_SIZE少16字节）
        in_file.read(reinterpret_cast<char*>(original_block.data()), BLOCK_SIZE - IV_SIZE);
        size_t bytes_read = in_file.gcount();
        if (bytes_read == 0)
            break;

        // 准备加密缓冲区（前16字节留给随机值）
        memcpy(processing_block.data() + IV_SIZE, original_block.data(), bytes_read);

        // 加密计时
        auto start_encrypt = std::chrono::high_resolution_clock::now();
        if (!encrypt(ik, processing_block.data(), bytes_read + IV_SIZE)) {
            std::cerr << "Encryption failed at block " << blocks_processed << std::endl;
            verification_passed = false;
            break;
        }
        auto end_encrypt = std::chrono::high_resolution_clock::now();
        total_encrypt_time +=
            std::chrono::duration_cast<std::chrono::microseconds>(end_encrypt - start_encrypt);

        // 解密计时（使用加密后的数据）
        auto start_decrypt = std::chrono::high_resolution_clock::now();
        if (!decrypt(ik, processing_block.data(), bytes_read + IV_SIZE)) {
            std::cerr << "Decryption failed at block " << blocks_processed << std::endl;
            verification_passed = false;
            break;
        }
        auto end_decrypt = std::chrono::high_resolution_clock::now();
        total_decrypt_time +=
            std::chrono::duration_cast<std::chrono::microseconds>(end_decrypt - start_decrypt);

        // 验证解密结果（比较解密后的数据部分）
        if (memcmp(original_block.data(), processing_block.data() + IV_SIZE, bytes_read) != 0) {
            std::cerr << "Data mismatch at block " << blocks_processed << std::endl;
            verification_passed = false;
            break;
        }

        // 写入解密后的数据（只写入有效数据部分）
        decrypted_file.write(reinterpret_cast<const char*>(processing_block.data() + IV_SIZE), bytes_read);

        total_bytes += bytes_read;
        blocks_processed++;
    }

    // 关闭文件
    in_file.close();
    decrypted_file.close();

    // 计算吞吐量（只计算有效数据部分）
    double encrypt_throughput =
        (double)total_bytes / (1024 * 1024) / (total_encrypt_time.count() / 1000000.0);
    double decrypt_throughput =
        (double)total_bytes / (1024 * 1024) / (total_decrypt_time.count() / 1000000.0);

    // 输出结果
    std::cout << "\nOptimized Block Encryption/Decryption Test\n";
    std::cout << "=========================================\n";
    std::cout << "File size: " << file_size << " bytes (" << (double)file_size / (1024 * 1024) << " MB)\n";
    std::cout << "Effective block size: " << (BLOCK_SIZE - IV_SIZE) << " bytes\n";
    std::cout << "IV size: " << IV_SIZE << " bytes\n";
    std::cout << "Blocks processed: " << blocks_processed << "\n";
    std::cout << "Total encryption time: " << total_encrypt_time.count() << " μs (" << encrypt_throughput
              << " MB/s)\n";
    std::cout << "Total decryption time: " << total_decrypt_time.count() << " μs (" << decrypt_throughput
              << " MB/s)\n";
    std::cout << "Decrypted file: " << decrypted_path << "\n";
    std::cout << "Data verification: " << (verification_passed ? "PASSED" : "FAILED") << "\n";
    std::cout << "=========================================\n";

    // 如果验证失败，删除可能不正确的解密文件
    if (!verification_passed) {
        fs::remove(decrypted_path);
    }
}

void base_test()
{
    std::string key = "sss";
    uint8_t ik[32]{};
    hash(key.c_str(), ik);

    int offset = 16;
    char* msg = new char[256]{};
    memset(msg, 0, 256);
    auto len = std::snprintf(msg + offset, 256 - offset, "%s", "hello world");
    std::cout << encrypt(ik, (uint8_t*)msg, len + offset) << std::endl;
    std::cout << msg + offset << std::endl;

    uint8_t ik2[32]{};
    hash(key.c_str(), ik2);

    std::cout << decrypt(ik2, (uint8_t*)msg, len + offset) << std::endl;

    std::cout << msg + offset << std::endl;
    delete[] msg;
}

int main()
{
    test_speed("D:\\1.iso");

    return 0;
}