/*******************************************************************************
 * Copyright (c) 2025, Jan Koester jan.koester@gmx.net
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 * Redistributions of source code must retain the above copyright
 *      notice, this list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright
 *      notice, this list of conditions and the following disclaimer in the
 *      documentation and/or other materials provided with the distribution.
 * Neither the name of the <organization> nor the
 *      names of its contributors may be used to endorse or promote products
 *      derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *******************************************************************************/

#include <iostream>
#include <vector>
#include <sstream>
#include <iomanip>
#include <stdexcept>
#include <algorithm>
#include <numeric>

#include "../src/crypto/des.h"

// Use the namespace for convenience in the test file
using namespace netplus;

// --- Helper Functions for Test Environment ---

uint8_t hexToByte(const std::string& hex) {
    if (hex.length() != 2) {
        throw std::invalid_argument("Input string must be two characters for hex conversion.");
    }
    return static_cast<uint8_t>(std::stoul(hex, nullptr, 16));
}

// Updated to use netplus::block64
block64 hexDecode(const std::string& hexString) {
    if (hexString.length() % 2 != 0) {
        throw std::invalid_argument("Hex string must have an even length.");
    }

    block64 bytes;
    for (size_t i = 0; i < hexString.length(); i += 2) {
        bytes.push_back(hexToByte(hexString.substr(i, 2)));
    }
    return bytes;
}

// Updated to use netplus::block64
std::string hexEncode(const block64& bytes) {
    std::stringstream ss;
    ss << std::hex << std::uppercase << std::setfill('0');
    for (uint8_t byte : bytes) {
        ss << std::setw(2) << static_cast<int>(byte);
    }
    return ss.str();
}

// --- Simple Assertion Macro for CTest ---
#define CTEST_ASSERT(condition, message) \
    if (!(condition)) { \
        std::cerr << "TEST FAILED: " << __FILE__ << ":" << __LINE__ << " -> " << message << std::endl; \
        return 1; \
    }

// --- Test Implementation Functions ---

int test_key_size_validation() {
    std::cout << "Running test_key_size_validation..." << std::endl;

    // Test 1: Invalid Key Size (7 bytes) - updated to use block64
    block64 invalidKey(7);
    bool exceptionCaught = false;

    try {
        netplus::des des(invalidKey); // netplus::DES
    } catch (const std::invalid_argument& e) {
        exceptionCaught = true;
    } catch (...) {
        // Catch other exceptions if they occur
    }

    CTEST_ASSERT(exceptionCaught, "DES constructor failed to throw std::invalid_argument for 7-byte key.")

    // Test 2: Valid Key Size (8 bytes) - updated to use block64
    block64 validKey(8);
    exceptionCaught = false;

    try {
        netplus::des des(validKey); // netplus::DES
    } catch (const std::invalid_argument& e) {
        exceptionCaught = true;
    } catch (...) {
        // Catch other exceptions if they occur
    }

    CTEST_ASSERT(!exceptionCaught, "DES constructor threw an exception for a valid 8-byte key.")

    return 0; // Success
}

int test_encryption_roundtrip() {
    std::cout << "Running test_encryption_roundtrip..." << std::endl;

    const std::string KEY_HEX = "133457799BBCDFF1";
    const std::string PLAINTEXT_HEX = "0123456789ABCDEF";

    // Updated to use block64
    block64 key = hexDecode(KEY_HEX);
    block64 plaintext = hexDecode(PLAINTEXT_HEX);

    try {
        netplus::des des(key); // netplus::DES

        block64 ciphertext = des.encrypt(plaintext);
        CTEST_ASSERT(ciphertext.size() == 8, "Ciphertext size is not 8 bytes.")

        block64 decryptedPlaintext = des.decrypt(ciphertext);

        std::string actualHex = hexEncode(decryptedPlaintext);
        std::string expectedHex = PLAINTEXT_HEX;

        CTEST_ASSERT(actualHex == expectedHex,
            "Roundtrip failed: Expected " + expectedHex + ", Got " + actualHex)

    } catch (const std::exception& e) {
        std::string message = "Exception during roundtrip test: ";
        message += e.what();
        CTEST_ASSERT(false, message)
    }

    return 0; // Success
}

int test_standard_vector_check() {
    std::cout << "Running test_standard_vector_check..." << std::endl;

    const std::string KEY_HEX = "133457799BBCDFF1";
    const std::string PLAINTEXT_HEX = "0123456789ABCDEF";
    const std::string EXPECTED_CIPHERTEXT_HEX = "85E813540F0AB405"; // NIST SP 800-67 Rev. 2

    // Updated to use block64
    block64 key = hexDecode(KEY_HEX);
    block64 plaintext = hexDecode(PLAINTEXT_HEX);

    try {
        netplus::des des(key); // netplus::DES
        block64 actualCiphertext = des.encrypt(plaintext);

        std::string actualHex = hexEncode(actualCiphertext);

        // This should now pass with the full DES implementation.
        CTEST_ASSERT(actualHex == EXPECTED_CIPHERTEXT_HEX,
            "Standard vector check FAILED. Expected: " + EXPECTED_CIPHERTEXT_HEX + ", Got: " + actualHex)

    } catch (const std::exception& e) {
        std::string message = "Exception during standard vector test: ";
        message += e.what();
        CTEST_ASSERT(false, message)
    }

    return 0; // Success
}

static std::string hex8(const netplus::block64& b) {
    std::ostringstream oss;
    oss << std::hex << std::uppercase << std::setfill('0');
    for (auto x : b) oss << std::setw(2) << (int)x;
    return oss.str();
}

static netplus::block64 E(const netplus::des& d, const netplus::block64& x) { return d.encrypt(x); }
static netplus::block64 D(const netplus::des& d, const netplus::block64& x) { return d.decrypt(x); }

static void dump_3des_variants(const netplus::block64& P,
                               const netplus::block64& K1,
                               const netplus::block64& K2,
                               const netplus::block64& K3,
                               const std::string& expectedHex)
{
    struct Key { const char* name; netplus::block64 k; };
    std::array<Key,3> keys{{ {"K1",K1}, {"K2",K2}, {"K3",K3} }};

    std::array<std::array<int,3>,6> perms{{
        {{0,1,2}}, {{0,2,1}}, {{1,0,2}},
        {{1,2,0}}, {{2,0,1}}, {{2,1,0}}
    }};

    for (auto p : perms) {
        netplus::des A(keys[p[0]].k), B(keys[p[1]].k), C(keys[p[2]].k);

        // A,B,C are the 1st,2nd,3rd stages (left-to-right)
        auto ede = E(C, D(B, E(A, P)));  // E D E
        auto eee = E(C, E(B, E(A, P)));  // E E E

        auto edeHex = hex8(ede);
        auto eeeHex = hex8(eee);

        std::cout << keys[p[0]].name << "-" << keys[p[1]].name << "-" << keys[p[2]].name
                  << "  EDE=" << edeHex << (edeHex == expectedHex ? "  <-- MATCH" : "")
                  << "  EEE=" << eeeHex << (eeeHex == expectedHex ? "  <-- MATCH" : "")
                  << "\n";
    }
}

/**
 * @brief Test standard Triple DES EDE (3-Key) encryption and decryption roundtrip.
 * * Uses the same key for all three DES stages (K1=K2=K3) for simplicity in roundtrip test,
 * though a true 3DES application would use three unique keys.
 */
int test_3des_encryption_roundtrip() {
    std::cout << "Running test_3des_encryption_roundtrip..." << std::endl;

    // A single 8-byte key, used three times (K1=K2=K3)
    const std::string KEY_HEX = "133457799BBCDFF1";
    const std::string PLAINTEXT_HEX = "0123456789ABCDEF";

    block64 key = hexDecode(KEY_HEX);
    block64 plaintext = hexDecode(PLAINTEXT_HEX);

    try {
        // Initialize 3DES with K1, K2, K3 (all equal to 'key' for this test)
        netplus::trippledes tdes(key, key, key);

        block64 ciphertext = tdes.encrypt(plaintext);
        CTEST_ASSERT(ciphertext.size() == 8, "3DES Ciphertext size is not 8 bytes.")

        block64 decryptedPlaintext = tdes.decrypt(ciphertext);

        std::string actualHex = hexEncode(decryptedPlaintext);
        std::string expectedHex = PLAINTEXT_HEX;

        CTEST_ASSERT(actualHex == expectedHex,
            "3DES Roundtrip failed: Expected " + expectedHex + ", Got " + actualHex)

    } catch (const std::exception& e) {
        std::string message = "Exception during 3DES roundtrip test: ";
        message += e.what();
        CTEST_ASSERT(false, message)
    }

    return 0; // Success
}

/**
 * @brief Known-answer test for 3-key Triple DES (TDEA / EDE3) in ECB mode.
 *
 * Source: NIST SP 800-67 (original), Appendix B.1 (Example of TDEA Forward Cipher Operations - ECB).
 * Keys (Keying Option 1 / 3-key):
 *   K1 = 0123456789ABCDEF
 *   K2 = 23456789ABCDEF01
 *   K3 = 456789ABCDEF0123
 *
 * Test block:
 *   P1 = 5468652071756663  ("The quic")
 *   C1 = A826FD8CE53B855F
 */
int test_3des_standard_vector_check() {
    std::cout << "Running test_3des_standard_vector_check (EDE3 KAT)..." << std::endl;

    const std::string K1_HEX = "0123456789ABCDEF";
    const std::string K2_HEX = "23456789ABCDEF01";
    const std::string K3_HEX = "456789ABCDEF0123";

    const std::string PLAINTEXT_HEX          = "5468652071756663"; // "The quic"
    const std::string EXPECTED_CIPHERTEXT_HEX = "A826FD8CE53B855F";

    block64 key1 = hexDecode(K1_HEX);
    block64 key2 = hexDecode(K2_HEX);
    block64 key3 = hexDecode(K3_HEX);
    block64 plaintext = hexDecode(PLAINTEXT_HEX);

    try {
        netplus::trippledes tdes(key1, key2, key3);

        // --- Encrypt KAT ---
        block64 ct = tdes.encrypt(plaintext);
        std::string ctHex = hexEncode(ct);

        CTEST_ASSERT(ctHex == EXPECTED_CIPHERTEXT_HEX,
            "3DES KAT ENCRYPT FAILED. Expected: " + EXPECTED_CIPHERTEXT_HEX + ", Got: " + ctHex)

        // --- Decrypt KAT ---
        block64 pt2 = tdes.decrypt(ct);
        std::string ptHex = hexEncode(pt2);

        CTEST_ASSERT(ptHex == PLAINTEXT_HEX,
            "3DES KAT DECRYPT FAILED. Expected: " + PLAINTEXT_HEX + ", Got: " + ptHex)

    } catch (const std::exception& e) {
        dump_3des_variants(plaintext,key1,key2,key3,EXPECTED_CIPHERTEXT_HEX);
        netplus::des d1(key1), d2(key2), d3(key3);
        auto ref = d1.encrypt(d2.decrypt(d3.encrypt(plaintext)));
        std::cout << "REF(K1-K2-K3 EDE)=" << hexEncode(ref) << "\n";
        CTEST_ASSERT(false, std::string("Exception during 3DES KAT: ") + e.what())
    }

    return 0;
}


// --- Main Entry Point for CTest ---
int main() {
    int result = 0;

    // DES tests
    result |= test_key_size_validation();
    result |= test_encryption_roundtrip();
    result |= test_standard_vector_check();

    // 3DES tests
    result |= test_3des_encryption_roundtrip();
    result |= test_3des_standard_vector_check();

    if (result != 0) {
        std::cerr << "\nOne or more unit tests failed." << std::endl;
        return 1;
    }

    std::cout << "\nAll unit tests passed." << std::endl;
    return 0; // Success
}

