#include <iostream>
#include <iomanip>
#include <vector>
#include <string>
#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 void dump_hex(const std::vector<uint8_t>& v, const char* label) {
    std::cerr << label << " (" << v.size() << " bytes): ";
    for (auto b : v)
        std::cerr << std::hex << std::setw(2) << std::setfill('0') << int(b);
    std::cerr << std::dec << "\n";
}

static void dump_hex_raw(const uint8_t* p, size_t n, const char* label) {
    std::cerr << label << " (" << n << " bytes): ";
    for (size_t i=0;i<n;i++)
        std::cerr << std::hex << std::setw(2) << std::setfill('0') << int(p[i]);
    std::cerr << std::dec << "\n";
}

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 int first_diff(const std::vector<uint8_t>& a, const std::vector<uint8_t>& b) {
    if (a.size() != b.size()) return -2;
    for (size_t i=0;i<a.size();i++) if (a[i] != b[i]) return (int)i;
    return -1;
}

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

    auto key = hex_to_vec("000102030405060708090a0b0c0d0e0f");
    auto pt  = hex_to_vec("00112233445566778899aabbccddeeff");
    auto exp = hex_to_vec("69c4e0d86a7b0430d8cdb78070b4c55a");

    netplus::aes128 A(key);

    auto ct = A.encrypt(pt);
    auto dec = A.decrypt(ct);

    dump_hex(key, "[ECB] key");
    dump_hex(pt,  "[ECB] pt ");
    dump_hex(ct,  "[ECB] ct ");
    dump_hex(exp, "[ECB] exp");
    dump_hex(dec, "[ECB] dec");

    if (!eq(ct, exp) || !eq(dec, pt)) {
        std::cerr << "❌ ECB mismatch\n";
        return false;
    }
    std::cout << "✅ AES block encrypt/decrypt PASS\n";
    return true;
}

static bool test_cbc() {
    std::cout << "\n==== AES-CBC encryptCBC/decryptCBC tests ====\n";

    auto key = hex_to_vec("2b7e151628aed2a6abf7158809cf4f3c");
    auto iv  = hex_to_vec("000102030405060708090a0b0c0d0e0f");

    auto pt  = hex_to_vec(
        "6bc1bee22e409f96e93d7e117393172a"
        "ae2d8a571e03ac9c9eb76fac45af8e51"
    );

    auto exp_ct = hex_to_vec(
        "7649abac8119b246cee98e9b12e9197d"
        "5086cb9b507219ee95db113a917678b2"
    );

    netplus::aes128 A(key);

    auto ct = A.encryptCBC(pt, iv);
    auto dec = A.decryptCBC(ct, iv);

    dump_hex(key, "[CBC] key");
    dump_hex(iv,  "[CBC] iv ");
    dump_hex(pt,  "[CBC] pt ");
    dump_hex(ct,  "[CBC] ct ");
    dump_hex(dec, "[CBC] dec");

    if (!eq(ct, exp_ct) || !eq(dec, pt)) {
        std::cerr << "❌ CBC mismatch\n";
        return false;
    }
    std::cout << "✅ CBC encryptCBC/decryptCBC PASS\n";
    return true;
}

static bool test_gcm_vector(const char* name,
                            const std::vector<uint8_t>& key,
                            const std::vector<uint8_t>& iv,
                            const std::vector<uint8_t>& aad,
                            const std::vector<uint8_t>& pt,
                            const std::vector<uint8_t>& exp_ct,
                            const std::vector<uint8_t>& exp_tag)
{
    std::cout << "\n==== AES-GCM " << name << " ====\n";

    netplus::aes128 A(key);

    std::vector<uint8_t> ct(pt.size(), 0);
    uint8_t tag[16] = {0};

    bool ok = A.aes_gcm_encrypt(iv.data(),
                               aad.data(), aad.size(),
                               pt.data(), pt.size(),
                               ct.data(),
                               tag);
    if (!ok) {
        std::cerr << "❌ encrypt returned false\n";
        return false;
    }

    std::vector<uint8_t> got_tag(tag, tag+16);

    dump_hex(key, "[GCM] key");
    dump_hex(iv,  "[GCM] iv ");
    dump_hex(aad, "[GCM] aad");
    dump_hex(pt,  "[GCM] pt ");
    dump_hex(ct,  "[GCM] ct ");
    dump_hex(got_tag, "[GCM] tag");

    dump_hex(exp_ct,  "[GCM] exp_ct ");
    dump_hex(exp_tag, "[GCM] exp_tag");

    if (!eq(ct, exp_ct)) {
        int d = first_diff(ct, exp_ct);
        std::cerr << "❌ Ciphertext mismatch at byte " << d << "\n";
        return false;
    }

    if (!eq(got_tag, exp_tag)) {
        int d = first_diff(got_tag, exp_tag);
        std::cerr << "❌ Tag mismatch\n";
        if (d >= 0) {
            std::cerr << "Mismatch detail:\n  at byte "
                      << d << " got=" << std::hex << int(got_tag[d])
                      << " exp=" << int(exp_tag[d]) << std::dec << "\n";
        }
        return false;
    }

    // decrypt ok
    std::vector<uint8_t> dec(pt.size(), 0);
    bool ok2 = A.aes_gcm_decrypt(iv.data(),
                                aad.data(), aad.size(),
                                ct.data(), ct.size(),
                                exp_tag.data(),
                                dec.data());

    dump_hex(dec, "[GCM] dec");

    if (!ok2 || !eq(dec, pt)) {
        std::cerr << "❌ decrypt mismatch or tag verify fail\n";
        return false;
    }

    // negative tag test
    uint8_t badtag[16];
    std::memcpy(badtag, exp_tag.data(), 16);
    badtag[0] ^= 0x01;

    std::vector<uint8_t> dummy(pt.size(), 0);
    bool ok3 = A.aes_gcm_decrypt(iv.data(),
                                aad.data(), aad.size(),
                                ct.data(), ct.size(),
                                badtag,
                                dummy.data());
    if (ok3) {
        std::cerr << "❌ negative tag test FAILED (should reject)\n";
        return false;
    }

    std::cout << "✅ " << name << " PASS (incl. negative tag test)\n";
    return true;
}

static bool test_gcm() {
    std::cout << "\n==== AES-GCM known test vectors ====\n";

    // Vector 1
    {
        std::vector<uint8_t> key(16, 0x00);
        std::vector<uint8_t> iv(12, 0x00);
        std::vector<uint8_t> aad;
        std::vector<uint8_t> pt;
        std::vector<uint8_t> exp_ct;
        std::vector<uint8_t> exp_tag = hex_to_vec("58e2fccefa7e3061367f1d57a4e7455a");

        if (!test_gcm_vector("vector #1", key, iv, aad, pt, exp_ct, exp_tag))
            return false;
    }

    // Vector 2 (NIST Test Case 2: 64 bytes PT, no AAD)
    {
        std::vector<uint8_t> key = hex_to_vec("feffe9928665731c6d6a8f9467308308");
        std::vector<uint8_t> iv  = hex_to_vec("cafebabefacedbaddecaf888");
        std::vector<uint8_t> aad;
        std::vector<uint8_t> pt  = hex_to_vec(
            "d9313225f88406e5a55909c5aff5269a"
            "86a7a9531534f7da2e4c303d8a318a72"
            "1c3c0c95956809532fcf0e2449a6b525"
            "b16aedf5aa0de657ba637b391aafd255"
        );
        std::vector<uint8_t> exp_ct = hex_to_vec(
            "42831ec2217774244b7221b784d0d49c"
            "e3aa212f2c02a4e035c17e2329aca12e"
            "21d514b25466931c7d8f6a5aac84aa05"
            "1ba30b396a0aac973d58e091473f5985"
        );
        std::vector<uint8_t> exp_tag = hex_to_vec("4d5c2af327cd64a62cf35abd2ba6fab4");

        if (!test_gcm_vector("vector #2", key, iv, aad, pt, exp_ct, exp_tag))
            return false;
    }

    return true;
}

static bool test_aes256() {
    std::cout << "\n==== AES-256 single-block encrypt/decrypt tests ====\n";

    // Test vector from NIST (Appendix C.1 from FIPS 197, but with 256-bit key)
    auto key = hex_to_vec(
        "000102030405060708090a0b0c0d0e0f"
        "101112131415161718191a1b1c1d1e1f"
    );
    auto pt  = hex_to_vec("00112233445566778899aabbccddeeff");
    
    try {
        netplus::aes256 A(key);
        auto ct = A.encrypt(pt);
        auto dec = A.decrypt(ct);

        dump_hex(key, "[AES-256] key");
        dump_hex(pt,  "[AES-256] pt ");
        dump_hex(ct,  "[AES-256] ct ");
        dump_hex(dec, "[AES-256] dec");

        if (!eq(pt, dec)) {
            std::cerr << "❌ AES-256 decrypt mismatch at byte " << first_diff(pt, dec) << "\n";
            return false;
        }
        std::cout << "✅ AES-256 block encrypt/decrypt PASS\n";

        // Test CBC mode
        std::cout << "\n==== AES-256 CBC mode test ====\n";
        auto key256 = hex_to_vec(
            "603deb1015ca71be2b73aef0857d7781"
            "1f352c073b6108d72d9810a30914dff4"
        );
        auto iv = hex_to_vec("000102030405060708090a0b0c0d0e0f");
        auto pt_cbc = hex_to_vec(
            "6bc1bee22e409f96e93d7e117393172a"
            "ae2d8a571e03ac9c9eb76fac45af8e51"
        );

        netplus::aes256 A2(key256);
        auto ct_cbc = A2.encryptCBC(pt_cbc, iv);
        auto dec_cbc = A2.decryptCBC(ct_cbc, iv);

        dump_hex(key256, "[AES-256 CBC] key");
        dump_hex(iv,     "[AES-256 CBC] iv ");
        dump_hex(pt_cbc, "[AES-256 CBC] pt ");
        dump_hex(ct_cbc, "[AES-256 CBC] ct ");

        if (!eq(pt_cbc, dec_cbc)) {
            std::cerr << "❌ AES-256 CBC decrypt mismatch at byte " << first_diff(pt_cbc, dec_cbc) << "\n";
            return false;
        }
        std::cout << "✅ AES-256 CBC PASS\n";

        // Test GCM mode
        std::cout << "\n==== AES-256 GCM mode test ====\n";
        auto key256_gcm = hex_to_vec(
            "0100000000000000000000000000000000000000000000000000000000000000"
        );
        auto iv_gcm = hex_to_vec("000000000000000000000000");
        auto pt_gcm = hex_to_vec("");
        auto exp_ct_gcm = hex_to_vec("");
        auto exp_tag_gcm = hex_to_vec("0dbbc8e9426965efb43028db8f630b62");

        netplus::aes256 A3(key256_gcm);
        uint8_t tag[16];
        std::vector<uint8_t> ct_gcm(pt_gcm.size());

        if (!A3.aes_gcm_encrypt(
            (const uint8_t*)iv_gcm.data(),
            nullptr, 0,
            pt_gcm.data(), pt_gcm.size(),
            ct_gcm.data(),
            tag
        )) {
            std::cerr << "❌ AES-256 GCM encrypt failed\n";
            return false;
        }

        dump_hex(key256_gcm, "[AES-256 GCM] key");
        dump_hex(iv_gcm,     "[AES-256 GCM] iv ");
        dump_hex_raw(tag, 16, "[AES-256 GCM] tag");
        dump_hex(exp_tag_gcm, "[AES-256 GCM] exp_tag");

        if (!eq(exp_tag_gcm, std::vector<uint8_t>(tag, tag+16))) {
            std::cerr << "❌ AES-256 GCM tag mismatch\n";
            return false;
        }

        uint8_t pt_out[1] = {0};
        if (!A3.aes_gcm_decrypt(
            (const uint8_t*)iv_gcm.data(),
            nullptr, 0,
            ct_gcm.data(), ct_gcm.size(),
            tag,
            pt_out
        )) {
            std::cerr << "❌ AES-256 GCM decrypt failed (tag verification)\n";
            return false;
        }

        std::cout << "✅ AES-256 GCM PASS\n";

        return true;
    } catch (const std::exception& e) {
        std::cerr << "❌ AES-256 exception: " << e.what() << "\n";
        return false;
    }
}

int main() {
    int fail = 0;

    if (!test_ecb()) fail++;
    if (!test_cbc()) fail++;
    if (!test_gcm()) fail++;
    if (!test_aes256()) fail++;

    if (fail) {
        std::cout << "\n❌ AES FULL UNIT TEST FAILED (" << fail << " group(s) failed)\n";
        return 1;
    }

    std::cout << "\n✅ All AES unit tests PASSED\n";
    return 0;
}
