#include <tinycrypt/aes.h>
#include <tinycrypt/ccm_mode.h>
#include <tinycrypt/constants.h>
#include <string.h>
#include <stdio.h>
#include "../lib/debug_config.h"

#define SC_TAG_SIZE 8
#define SC_SESSION_KEY_SIZE 16

int main() {
    debug_config_init();
    debug_set_level(DEBUG_LEVEL_TRACE);
    debug_set_categories(DEBUG_CATEGORY_ALL);
    
    DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "=== Basic Crypto Test ===");
    
    struct tc_aes_key_sched_struct sched;
    uint8_t test_key[SC_SESSION_KEY_SIZE] = {0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 
                                              0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x10};
    
    if (tc_aes128_set_encrypt_key(&sched, test_key) != TC_CRYPTO_SUCCESS) {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "AES key setup failed");
        return 1;
    }
    
    DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "AES key setup successful");
    
    struct tc_ccm_mode_struct ccm_state;
    TCCcmMode_t c = &ccm_state;
    uint8_t nonce[13] = {0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xAA, 0xBB, 0xCC, 0xDD};
    
    if (tc_ccm_config(c, &sched, nonce, 13, SC_TAG_SIZE) != TC_CRYPTO_SUCCESS) {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "CCM config failed");
        return 1;
    }
    
    DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "CCM config successful");
    
    uint8_t plaintext[] = "Hello, World!";
    uint8_t ciphertext[64];
    size_t plaintext_len = sizeof(plaintext) - 1;
    size_t ciphertext_len = plaintext_len + SC_TAG_SIZE;
    
    if (tc_ccm_generation_encryption(ciphertext, ciphertext_len, NULL, 0, plaintext, plaintext_len, c) != TC_CRYPTO_SUCCESS) {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "Encryption failed");
        return 1;
    }
    
    DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "Encryption successful");
    
    uint8_t decrypted[64];
    if (tc_ccm_decryption_verification(decrypted, plaintext_len, NULL, 0, ciphertext, ciphertext_len, c) != TC_CRYPTO_SUCCESS) {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "Decryption failed");
        return 1;
    }
    
    DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "Decryption successful");
    
    if (memcmp(plaintext, decrypted, plaintext_len) == 0) {
        DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "Crypto test PASSED");
        return 0;
    } else {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "Data mismatch");
        return 1;
    }
}