#include <Arduino.h>

#define XTEA_ROUNDS 32
#define BLOCK_SIZE 8  // XTEA operates on 64-bit blocks
//#define MESSAGE_SIZE 32  // Message before padding
#define MESSAGE_SIZE 256

size_t NUM_INTS = MESSAGE_SIZE / BLOCK_SIZE * 2;
size_t NUM_BLKS = MESSAGE_SIZE / BLOCK_SIZE;


void zero_buffer(uint32_t* a) {
  for (size_t i = 0; i < NUM_INTS; i++) {
    a[i] = 0;
  }
} 


void convert_uint32_to_bytes(const uint32_t* input, byte* output, size_t blocks) {
    zero_buffer((uint32_t*) output);
    for (size_t i = 0; i < blocks; i++) {
        output[i * 4] = (input[i] >> 24) & 0xFF;  // Extract highest byte
        output[i * 4 + 1] = (input[i] >> 16) & 0xFF;
        output[i * 4 + 2] = (input[i] >> 8) & 0xFF;
        output[i * 4 + 3] = input[i] & 0xFF;      // Extract lowest byte
    }
}


// Function to correctly convert a byte array into a uint32_t array (big-endian format)
void convert_bytes_to_uint32(const byte* input, uint32_t* output, size_t blocks) {
    zero_buffer(output);
    for (size_t i = 0; i < blocks; i++) {
        output[i] = ((uint32_t)input[i * 4] << 24) | 
                    ((uint32_t)input[i * 4 + 1] << 16) | 
                    ((uint32_t)input[i * 4 + 2] << 8) | 
                    ((uint32_t)input[i * 4 + 3]);
    }
}


// XTEA encryption function
void xtea_encrypt(uint32_t v[2], const uint32_t key[4]) {
    uint32_t v0 = v[0], v1 = v[1], sum = 0, delta = 0x9E3779B9;
    for (int i = 0; i < XTEA_ROUNDS; i++) {
        v0 = (v0 + (((v1 << 4) ^ (v1 >> 5)) + v1 ^ (sum + key[sum & 3]))) & 0xFFFFFFFF;
        sum = (sum + delta) & 0xFFFFFFFF;
        v1 = (v1 + (((v0 << 4) ^ (v0 >> 5)) + v0 ^ (sum + key[(sum >> 11) & 3]))) & 0xFFFFFFFF;
    }
    v[0] = v0;
    v[1] = v1;
}


// XTEA decryption function
void xtea_decrypt(uint32_t v[2], const uint32_t key[4]) {
    uint32_t v0 = v[0], v1 = v[1], sum = XTEA_ROUNDS * 0x9E3779B9, delta = 0x9E3779B9;
    for (int i = 0; i < XTEA_ROUNDS; i++) {
        v1 = (v1 - (((v0 << 4) ^ (v0 >> 5)) + v0 ^ (sum + key[(sum >> 11) & 3]))) & 0xFFFFFFFF;
        sum = (sum - delta) & 0xFFFFFFFF;
        v0 = (v0 - (((v1 << 4) ^ (v1 >> 5)) + v1 ^ (sum + key[sum & 3]))) & 0xFFFFFFFF;
    }
    v[0] = v0;
    v[1] = v1;
}


// Encrypt message block by block
void encrypt_message(const uint32_t* msg, const uint32_t key[4], uint32_t* encrypted, size_t blocks) {
    for (size_t i = 0; i < blocks; i++) {
        uint32_t temp_block[2] = { msg[i * 2], msg[i * 2 + 1] };  // Copy original values
        xtea_encrypt(temp_block, key);  // Encrypt copy, not original
        encrypted[i * 2] = temp_block[0];
        encrypted[i * 2 + 1] = temp_block[1];
    }
}


// Decrypt message block by block
void decrypt_message(const uint32_t* encrypted, const uint32_t key[4], uint32_t* decrypted, size_t blocks) {
    for (size_t i = 0; i < blocks; i++) {
        uint32_t temp_block[2] = { encrypted[i * 2], encrypted[i * 2 + 1] };
        xtea_decrypt(temp_block, key);
        decrypted[i * 2] = temp_block[0];
        decrypted[i * 2 + 1] = temp_block[1];
    }
}


void print_hex(char* name, uint32_t* b) {
   Serial.println(name);
    for (size_t i = 0; i < NUM_INTS; i++) {
        Serial.printf("%08X ", b[i]);
    }
    Serial.println();
}


void print_bytes(char* name, byte* b) {
    Serial.println(name);
    Serial.println((char *) b);
}


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

    uint32_t key[4] = { 0x01234567, 0x89ABCDEF, 0xFEDCBA98, 0x76543210 };
   
    byte message[MESSAGE_SIZE] = "ABCDEFGH";  // Message as raw bytes
    print_bytes("original:", message);

    // Convert byte array into correct uint32_t format
    uint32_t formatted_message[NUM_INTS];
    convert_bytes_to_uint32(message, formatted_message, NUM_BLKS);
    print_hex("formatted:", formatted_message);

    uint32_t encrypted[NUM_INTS];
    uint32_t decrypted[NUM_INTS];

    encrypt_message(formatted_message, key, encrypted, NUM_BLKS);
    print_hex("encrypted:", encrypted);

    decrypt_message(encrypted, key, decrypted, NUM_BLKS);
    print_hex("decrypted:", decrypted);
  
    byte reconstructed_message[MESSAGE_SIZE];
    convert_uint32_to_bytes(decrypted, reconstructed_message, NUM_BLKS);
    print_bytes("reconstructed:", reconstructed_message);
}


void loop() {}