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

#define SC_TAG_SIZE 8  // This is valid: 8 bytes
#define SC_SESSION_KEY_SIZE 16

int main() {
    printf("=== Working Crypto Test ===\n");
    
    // Test AES key setup
    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};
    
    printf("Testing AES key setup...\n");
    int aes_result = tc_aes128_set_encrypt_key(&sched, test_key);
    printf("tc_aes128_set_encrypt_key returned: %d\n", aes_result);
    
    if (aes_result == TC_CRYPTO_SUCCESS) {
        printf("✓ AES key setup successful\n");
        
        // Test CCM config with CORRECT parameters
        struct tc_ccm_mode_struct ccm_state = {0};
        TCCcmMode_t c = &ccm_state;
        
        // CRITICAL: nonce MUST be exactly 13 bytes for CCM!
        uint8_t nonce[13] = {0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 
                             0x99, 0xAA, 0xBB, 0xCC, 0xDD};
        
        printf("\nTesting CCM config with CORRECT parameters...\n");
        printf("Nonce size: 13 (required by CCM standard)\n");
        printf("Tag size: %d (valid: 4,6,8,10,12,14,16)\n", SC_TAG_SIZE);
        printf("Using TCCcmMode_t pointer: %p\n", (void*)c);
        
        int ccm_result = tc_ccm_config(c, &sched, nonce, 13, SC_TAG_SIZE);
        printf("tc_ccm_config returned: %d\n", ccm_result);
        
        if (ccm_result == TC_CRYPTO_SUCCESS) {
            printf("✓ CCM config successful!\n");
            
            // Test simple encryption
            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;
            
            printf("\nTesting CCM encryption...\n");
            printf("Plaintext: '%s' (%zu bytes)\n", plaintext, plaintext_len);
            printf("Expected ciphertext: %zu bytes (data + tag)\n", ciphertext_len);
            
            int encrypt_result = tc_ccm_generation_encryption(ciphertext, ciphertext_len,
                                                             NULL, 0, /* no associated data */
                                                             plaintext, plaintext_len,
                                                             c);
            printf("tc_ccm_generation_encryption returned: %d\n", encrypt_result);
            
            if (encrypt_result == TC_CRYPTO_SUCCESS) {
                printf("🎉 CCM encryption successful!\n");
                printf("Ciphertext length: %zu bytes\n", ciphertext_len);
                
                // Test decryption
                uint8_t decrypted[64];
                size_t decrypted_len = plaintext_len;
                
                printf("\nTesting CCM decryption...\n");
                int decrypt_result = tc_ccm_decryption_verification(decrypted, decrypted_len,
                                                                   NULL, 0, /* no associated data */
                                                                   ciphertext, ciphertext_len,
                                                                   c);
                printf("tc_ccm_decryption_verification returned: %d\n", decrypt_result);
                
                if (decrypt_result == TC_CRYPTO_SUCCESS) {
                    printf("🎉 CCM decryption successful!\n");
                    printf("Decrypted: '%s'\n", decrypted);
                    
                    // Verify decrypted data
                    if (decrypted_len == plaintext_len && memcmp(plaintext, decrypted, plaintext_len) == 0) {
                        printf("✓ Decrypted data matches original!\n");
                        printf("\n🎉 CRYPTO WORKS! AES-CCM encryption/decryption successful!\n");
                        return 0;
                    } else {
                        printf("❌ Decrypted data doesn't match original\n");
                        return 1;
                    }
                } else {
                    printf("❌ CCM decryption failed\n");
                    return 1;
                }
            } else {
                printf("❌ CCM encryption failed\n");
                return 1;
            }
        } else {
            printf("❌ CCM config failed\n");
            return 1;
        }
    } else {
        printf("❌ AES key setup failed\n");
        return 1;
    }
}