DES_PAO

#include <iostream>
#include <fstream>
#include <cstdint>
#include <chrono>
#include <omp.h> // Include OpenMP header
#include "keys.h"
#include "permutation.h"
#include "sbox.h"

class TripleDES {
public:
    TripleDES(const uint64_t* key1, const uint64_t* key2, const uint64_t* key3) : key1(key1), key2(key2), key3(key3) {
        obtainRoundKeys(*key1, roundKeys1, 16);
        obtainRoundKeys(*key2, roundKeys2, 16);
        obtainRoundKeys(*key3, roundKeys3, 16);
    }

    void encryptFile(const char* inputFile, const char* outputFile);
    uint64_t generateNextCounterValue(uint64_t counter);
    uint64_t getMillisTimestamp();

private:
    const uint64_t* key1;
    const uint64_t* key2;
    const uint64_t* key3;
    unsigned long long roundKeys1[16], roundKeys2[16], roundKeys3[16];

    void encryptBlocks(const uint64_t* counter, uint64_t* ciphertextBlock);
};

void TripleDES::encryptFile(const char* inputFile, const char* outputFile) {
    std::ifstream inFile(inputFile, std::ios::binary);
    std::ofstream outFile(outputFile, std::ios::binary);

    if (!inFile.is_open() || !outFile.is_open()) {
        std::cerr << "Error opening files." << std::endl;
        return;
    }

    constexpr size_t blockSize = 8; // 64 bits
    uint64_t plaintextBlock[8], ciphertextBlock[8];
    uint64_t counter[8] = {0, 1, 2, 3, 4, 5, 6, 7}; // Initial counter values
    bool keepProcessing = true;

#pragma omp parallel // Start parallel region
    {
        while (keepProcessing) {
            // Read multiple blocks in a thread-safe manner
            bool readSuccessful = true;
#pragma omp critical
            {
                for (int i = 0; i < 8; ++i) {
                    readSuccessful &= inFile.read(reinterpret_cast<char*>(&plaintextBlock[i]), blockSize).good();
                }
            }

            if (!readSuccessful) {
                keepProcessing = false;
            }

            // Encrypt multiple blocks in parallel
            encryptBlocks(counter, ciphertextBlock);

            // Write the ciphertext blocks in a thread-safe manner
#pragma omp critical
            {
                if (keepProcessing) {
                    for (int i = 0; i < 8; ++i) {
                        outFile.write(reinterpret_cast<char*>(&ciphertextBlock[i]), blockSize);
                    }
                }
            }

            // Update the counter after processing the blocks
            for (int i = 0; i < 8; ++i) {
                counter[i] = generateNextCounterValue(counter[i]);
            }
        }
    } // End of parallel region

    inFile.close();
    outFile.close();
}

void TripleDES::encryptBlocks(const uint64_t* counter, uint64_t* ciphertextBlock) {
    // First encryption
#pragma omp simd
    for (int i = 0; i < 8; ++i) {
        uint64_t permutedBlock = performLongToLongPermutation(counter[i], initialPerm, initialPermSize);
        for (int j = 0; j < 16; ++j) {
            unsigned long long expandedBlock = performLongToLongPermutation(permutedBlock, expandRight, expandRightSize);
            unsigned long long xorResult = expandedBlock ^ roundKeys1[j];
            unsigned int sboxResult = performSboxCalculation(xorResult);
            uint64_t permutedSboxResult = performIntToLongPermutation(sboxResult, finalSboxPerm, finalSboxPermSize);
            permutedBlock ^= permutedSboxResult;
        }
        ciphertextBlock[i] = performLongToLongPermutation(permutedBlock, keyPermFinal, keyPermFinalSize);
    }

    // Second encryption
#pragma omp simd
    for (int i = 0; i < 8; ++i) {
        uint64_t permutedBlock = performLongToLongPermutation(ciphertextBlock[i], initialPerm, initialPermSize);
        for (int j = 0; j < 16; ++j) {
            unsigned long long expandedBlock = performLongToLongPermutation(permutedBlock, expandRight, expandRightSize);
            unsigned long long xorResult = expandedBlock ^ roundKeys2[j];
            unsigned int sboxResult = performSboxCalculation(xorResult);
            uint64_t permutedSboxResult = performIntToLongPermutation(sboxResult, finalSboxPerm, finalSboxPermSize);
            permutedBlock ^= permutedSboxResult;
        }
        ciphertextBlock[i] = performLongToLongPermutation(permutedBlock, keyPermFinal, keyPermFinalSize);
    }

    // Third encryption
#pragma omp simd
    for (int i = 0; i < 8; ++i) {
        uint64_t permutedBlock = performLongToLongPermutation(ciphertextBlock[i], initialPerm, initialPermSize);
        for (int j = 0; j < 16; ++j) {
            unsigned long long expandedBlock = performLongToLongPermutation(permutedBlock, expandRight, expandRightSize);
            unsigned long long xorResult = expandedBlock ^ roundKeys3[j];
            unsigned int sboxResult = performSboxCalculation(xorResult);
            uint64_t permutedSboxResult = performIntToLongPermutation(sboxResult, finalSboxPerm, finalSboxPermSize);
            permutedBlock ^= permutedSboxResult;
        }
        ciphertextBlock[i] = performLongToLongPermutation(permutedBlock, keyPermFinal, keyPermFinalSize);
    }
}

uint64_t TripleDES::generateNextCounterValue(uint64_t counter) {
    return counter + 8; // Increment by the number of blocks processed
}

uint64_t TripleDES::getMillisTimestamp() {
    auto now = std::chrono::system_clock::now();
    auto duration = now.time_since_epoch();
    auto millis = std::chrono::duration_cast<std::chrono::milliseconds>(duration).count();
    return millis;
}

int main() {
    const char* inputFile = "../data/input1MB.txt";
    const char* outputFile = "../data/output1MB_encrypted.txt";
    uint64_t key1[3] = {0xA4B92905A4B92905, 0x0000000000000000, 0xA4B92905A4B92905};
    uint64_t key2[3] = {0xC567A1F3C567A1F3, 0x0000000000000000, 0xC567A1F3C567A1F3};
    uint64_t key3[3] = {0xE3D17F253E3D17F2, 0x0000000000000000, 0xE3D17F253E3D17F2};

    TripleDES des(key1, key2, key3);

    auto begin = des.getMillisTimestamp();
    des.encryptFile(inputFile, outputFile);
    auto end = des.getMillisTimestamp();

    auto duration_ms = end - begin;
    std::cout << "Encryption completed in " << duration_ms << " ms." << std::endl;

    return 0;
}