#include <iostream>
#include <string>
#include <vector>
#include <cstring>
#include "../src/crypto/aes.h"

static std::vector<uint8_t> hex_to_vec(const char* hex) {
    std::vector<uint8_t> out;
    while (*hex) {
        while (*hex == ' ' || *hex == '\n' || *hex == '\t') hex++;
        if (!*hex) break;
        uint8_t hi = (uint8_t)std::stoi(std::string(hex, 1), nullptr, 16);
        uint8_t lo = (uint8_t)std::stoi(std::string(hex+1, 1), nullptr, 16);
        out.push_back((hi<<4) | lo);
        hex += 2;
    }
    return out;
}

static bool eq(const std::vector<uint8_t>& a, const std::vector<uint8_t>& b) {
    return a.size() == b.size() && std::memcmp(a.data(), b.data(), a.size()) == 0;
}

static bool test_cts() {
    std::cout << "\n==== AES-128 CTS encrypt/decrypt tests ====\n";

    auto key = hex_to_vec("000102030405060708090a0b0c0d0e0f");
    auto iv  = hex_to_vec("000102030405060708090a0b0c0d0e0f");
    // For CTS, the plaintext can be any length >= 16 bytes.
    // Let's test with a 32-byte plaintext.
    auto pt  = hex_to_vec("00112233445566778899aabbccddeeff"
                           "00112233445566778899aabbccddeeff");

    netplus::aes128 A(key);

    auto ct = A.encryptCTS(pt, iv);
    auto dec = A.decryptCTS(ct, iv);

    if (ct.size() != pt.size()) {
        std::cerr << "❌ CTS ciphertext size mismatch\n";
        return false;
    }

    if (!eq(dec, pt)) {
        std::cerr << "❌ CTS decrypt mismatch\n";
        return false;
    }

    std::cout << "✅ AES-128 CTS PASS\n";
    return true;
}

int main() {
    if (!test_cts()) {
        return 1;
    }
    std::cout << "\n✅ All AES-CTS unit tests PASSED\n";
    return 0;
}
