#include "AESLib.h"
#include "arduino_base64.hpp"
#include <Wire.h>
#include <LiquidCrystal_I2C.h>
#include <Keypad.h>

// Constants for LCD Display
LiquidCrystal_I2C lcd(0x27, 16, 2); // I2C address 0x27, 16 columns and 2 rows

// Define SDA and SCL pins for I2C communication
#define SDA_PIN 21
#define SCL_PIN 22

// Constants for Keypad
const byte ROWS = 4; // four rows
const byte COLS = 4; // four columns
char keys[ROWS][COLS] = {
    {'1', '2', '3', 'A'},
    {'4', '5', '6', 'B'},
    {'7', '8', '9', 'C'},
    {'*', '0', '#', 'D'}};
byte rowPins[ROWS] = {16, 17, 5, 18}; // connect to the row pinouts of the keypad
byte colPins[COLS] = {15, 2, 0, 4};  // connect to the column pinouts of the keypad
Keypad keypad = Keypad(makeKeymap(keys), rowPins, colPins, ROWS, COLS);

AESLib aesLib;

String inputText = ""; // Variable to store user input
String inputText2 = ""; // Variable to store user input

void setup()
{
    Serial.begin(115200);

    Wire.begin(SDA_PIN, SCL_PIN); // Initialize I2C with custom SDA and SCL pins

    lcd.begin(16, 2, LCD_5x8DOTS);
    lcd.backlight();

    Serial.println();
    Serial.println("--- Device Started ---");
}

void loop()
{
    char key = keypad.getKey();

    if (key)
    {
        // Key press detected
        if (key == '#')
        {
            // Encryption process when '#' key is pressed
            String encryptedText = encrypt(inputText);

            // Print original and encrypted text to Serial
            Serial.println("Original text: \"" + inputText + "\"");
            Serial.println("Encrypted text: \"" + encryptedText + "\"");

            // Display original and encrypted text on the LCD
            lcd.clear();
            lcd.setCursor(0, 0);
            lcd.print("Original: " + inputText);
            lcd.setCursor(0, 1);
            lcd.print("Encrypted: " + encryptedText);

            delay(2000);
            lcd.clear();

            // Reset inputText for the next entry
            inputText2 = inputText;
            inputText = "";
        }
        else if (key == '*')
        {
            // Decryption process when '*' key is pressed
            String decryptedText = decrypt(encrypt(inputText)); // Decrypt the previously encrypted text

            // Display the original text on the decrypted key field
            lcd.clear();
            lcd.setCursor(0, 0);
            lcd.print("Original: " + decryptedText);

            // Print decrypted text to Serial
            Serial.println("Decrypted text: \"" + inputText2 + "\"");

            lcd.setCursor(0, 1);
            lcd.print("Decrypted: " + decryptedText);

            delay(2000);
            lcd.clear();
        }
        else
        {
            // Accumulate entered letters
            inputText += key;

            // Display the entire inputText on the LCD
            lcd.clear();
            lcd.setCursor(0, 0);
            lcd.print("Input: " + inputText);
        }
    }
}

String encrypt(String inputText)
{
    // Convert input text to bytes
    int bytesInputLength = inputText.length() + 1;
    char bytesInput[bytesInputLength];
    inputText.toCharArray(bytesInput, bytesInputLength);

    // Initialize variables for encryption
    int outputLength = aesLib.get_cipher_length(bytesInputLength);
    byte bytesEncrypted[outputLength];

    byte aesKey[] = {23, 45, 56, 67, 67, 87, 98, 12, 32, 34, 45, 56, 67, 87, 65, 5};
    byte aesIv[] = {123, 43, 46, 89, 29, 187, 58, 213, 78, 50, 19, 106, 205, 1, 5, 7};

    // Set padding mode and convert char* to byte*
    aesLib.set_paddingmode((paddingMode)0);
    byte tempBuffer[bytesInputLength];
    memcpy(tempBuffer, bytesInput, bytesInputLength);

    // Encrypt using AES
    aesLib.encrypt(tempBuffer, bytesInputLength, bytesEncrypted, aesKey, 16, aesIv);

    // Encode encrypted bytes to base64 and return as a String
    char base64EncodedOutput[base64::encodeLength(outputLength)];
    base64::encode(bytesEncrypted, outputLength, base64EncodedOutput);

    return String(base64EncodedOutput);
}

String decrypt(String encryptedText)
{
    // Decode base64 to get original byte array length
    int originalBytesLength = base64::decodeLength(encryptedText.c_str());
    byte encryptedBytes[originalBytesLength];
    byte decryptedBytes[originalBytesLength];

    // Decode base64 to get encrypted byte array
    base64::decode(encryptedText.c_str(), encryptedBytes);

    byte aesKey[] = {23, 45, 56, 67, 67, 87, 98, 12, 32, 34, 45, 56, 67, 87, 65, 5};
    byte aesIv[] = {123, 43, 46, 89, 29, 187, 58, 213, 78, 50, 19, 106, 205, 1, 5, 7};

    // Set padding mode and decrypt using AES
    aesLib.set_paddingmode((paddingMode)0);
    aesLib.decrypt(encryptedBytes, originalBytesLength, decryptedBytes, aesKey, 16, aesIv);

    // Convert decrypted byte array to String
    String decryptedText = "";
    for (int i = 0; i < originalBytesLength; i++)
    {
        decryptedText += (char)decryptedBytes[i];
    }

    return decryptedText;
}