#include <assert.h>
#include <stdio.h>

#include "aes_internal.h"
#include "array.h"

#define TEST_DATA { 72, 101, 114, 101, 32, 105, 115, 32, 115, 111, 109, 101, 32, 116, 101, 120, 116, 32, 105, 110, 32, 97, 115, 99, 105, 105, 32, 115, 105, 110, 99, 101, 32, 116, 104, 101, 115, 101, 32, 97, 114, 101, 32, 97, 108, 108, 32, 98, 121, 116, 101, 115, 0 }
#define TEST_DATA_LENGTH 53
#define TEST_BLOCK { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, }
#define TEST_KEY { 130, 191, 5, 162, 175, 104, 200, 14, 32, 0, 97, 170, 10, 83, 159, 90, }

void aes_test_shift_block() {
    uint8_t original_data[BLOCK_LENGTH] = TEST_BLOCK;
    uint8_t test_data[BLOCK_LENGTH] = TEST_BLOCK;
    uint8_t expected_shift[BLOCK_LENGTH] = {
        0, 1, 2, 3, 5, 6, 7, 4, 10, 11, 8, 9, 15, 12, 13, 14,
    };

    aes__shift_rows(test_data, encryption);
    assert(arrays_eq(test_data, expected_shift, BLOCK_LENGTH));

    aes__shift_rows(test_data, decryption);
    assert(arrays_eq(test_data, original_data, BLOCK_LENGTH));
}

void aes_test_xtime() {
    assert(aes__xtime(84) == 168);
    assert(aes__xtime(255) == 229);
    assert(aes__xtime(0) == 0);
    assert(aes__xtime(240) == 251);
}

void aes_test_mix_words() {
    uint8_t test_data[BLOCK_LENGTH] = TEST_BLOCK;
    uint8_t original_data[BLOCK_LENGTH] = TEST_BLOCK;

    aes__mix_words(test_data, encryption);
    aes__mix_words(test_data, decryption);
    assert(arrays_eq(test_data, original_data, BLOCK_LENGTH));
}

void aes_test_expand_key() {
    uint8_t test_key[BLOCK_LENGTH] = TEST_KEY;
    uint8_t expected_round_keys[BLOCK_LENGTH * 11] = {
        130, 191, 5, 162, 175, 104, 200, 14, 32, 0, 97, 170, 10, 83, 159, 90, 110, 100, 187, 197, 193, 12, 115, 203, 225, 12, 18, 97, 235, 95, 141, 59, 163, 57, 89, 44, 98, 53, 42, 231, 131, 57, 56, 134, 104, 102, 181, 189, 148, 236, 35, 105, 246, 217, 9, 142, 117, 224, 49, 8, 29, 134, 132, 181, 216, 179, 246, 205, 46, 106, 255, 67, 91, 138, 206, 75, 70, 12, 74, 254, 54, 101, 77, 151, 24, 15, 178, 212, 67, 133, 124, 159, 5, 137, 54, 97, 177, 96, 162, 252, 169, 111, 16, 40, 234, 234, 108, 183, 239, 99, 90, 214, 10, 222, 84, 35, 163, 177, 68, 11, 73, 91, 40, 188, 166, 56, 114, 106, 141, 158, 86, 7, 46, 47, 18, 12, 103, 116, 58, 176, 193, 76, 72, 218, 191, 204, 1, 127, 145, 227, 19, 115, 246, 151, 41, 195, 55, 219, 97, 25, 48, 35, 213, 229, 161, 192, 198, 150, 87, 87, 239, 85, 96, 140, 142, 76
    };

    uint8_t* round_keys = aes__expand_key(test_key);
    int result = arrays_eq(round_keys, expected_round_keys, BLOCK_LENGTH);
    free(round_keys);

    assert(result);
}

void aes_test_block_cipher() {
    uint8_t test_key[BLOCK_LENGTH] = TEST_KEY;
    uint8_t* round_keys = aes__expand_key(test_key);

    uint8_t original_data[BLOCK_LENGTH] = TEST_BLOCK;
    uint8_t test_data[BLOCK_LENGTH] = TEST_BLOCK;

    aes__encrypt_block(test_data, round_keys);
    aes__decrypt_block(test_data, round_keys);

    free(round_keys);

    assert(arrays_eq(test_data, original_data, BLOCK_LENGTH));
}

void aes_test_padding_validation() {
    uint8_t data_valid[16] = {
        12, 63, 3, 76, 10, 92,
        10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
    };
    assert(aes__is_valid_padding(data_valid) == TRUE);

    uint8_t data_invalid[16] = {
        12, 63, 3, 76, 10, 92,
        10, 10, 10, 10, 10, 10, 10, 10, 10, 0,
    };
    assert(aes__is_valid_padding(data_invalid) == FALSE);

    uint8_t data_valid_all_pad[16] = {
        16, 16, 16, 16, 16, 16, 16, 16,
        16, 16, 16, 16, 16, 16, 16, 16,
    };
    assert(aes__is_valid_padding(data_valid_all_pad) == TRUE);
}

void aes_test_create_padding() {
    uint8_t data[24] = {
        12, 63,  3, 76, 10, 92, 34, 11,
        12, 63,  3, 76, 10, 92, 34, 11,

        12, 63,  3, 76, 10, 92, 34, 11,
    };

    uint8_t expected_padding[BLOCK_LENGTH] = {
        12, 63,  3, 76, 10, 92, 34, 11,
        8 ,  8,  8,  8,  8,  8,  8,  8,
    };

    aes_padded_data_t padded = aes_create_padded_data_container();
    int code = aes_pad_data(&padded, data, 24);
    assert(code == 0);
    assert(padded.data == data);
    assert(padded.data_length == 24);
    assert(padded.data_length_before_pad == 16);
    assert(padded.length == 32);
    assert(arrays_eq(padded.padded_block, expected_padding, BLOCK_LENGTH));

    uint8_t* data_ret = aes_drop_padded_data_container(&padded);
    assert(data_ret == data);
}

// TODO: Doesn't test padded block
void aes_test_cipher() {
    uint8_t original[TEST_DATA_LENGTH] = TEST_DATA;
    uint8_t data[TEST_DATA_LENGTH] = TEST_DATA;
    uint8_t key[BLOCK_LENGTH] = TEST_KEY;

    aes_padded_data_t padded = aes_create_padded_data_container();
    int code = aes_pad_data(&padded, data, TEST_DATA_LENGTH);
    assert(code == 0);

    // ECB

    aes_ecb(encryption, &padded, key);
    assert(!arrays_eq(padded.data, original, padded.data_length_before_pad));

    aes_ecb(decryption, &padded, key);
    assert(arrays_eq(padded.data, original, padded.data_length_before_pad));

    // CBC

    aes_cbc(encryption, &padded, key);
    assert(!arrays_eq(padded.data, original, padded.data_length_before_pad));

    aes_cbc(decryption, &padded, key);
    assert(arrays_eq(padded.data, original, padded.data_length_before_pad));

    aes_drop_padded_data_container(&padded);
}