#include "../src/secure_channel.h"
#include <tinycrypt/aes.h>
#include <tinycrypt/ccm_mode.h>
#include <tinycrypt/constants.h>
#include <stdint.h>
#include <string.h>
#include <stdio.h>

#define SC_NONCE_SIZE 8
#define SC_TAG_SIZE 8
#define SC_SESSION_KEY_SIZE 16
#define SC_CRC32_SIZE 4

int main() {
    printf("=== Step-by-Step Debug ===\n");
    
    // 1. Setup context
    struct SC_MYKEYS keys;
    for (int i = 0; i < 32; i++) {
        keys.private_key[i] = i;
        keys.public_key[i] = i + 64;
    }
    
    sc_context_t ctx;
    sc_init_ctx(&ctx, &keys);
    
    // Manually setup for encryption
    memcpy(ctx.peer_public_key, keys.public_key, 64);
    ctx.peer_key_set = 1;
    ctx.session_ready = 1;
    
    uint8_t test_session_key[SC_SESSION_KEY_SIZE];
    for (int i = 0; i < SC_SESSION_KEY_SIZE; i++) {
        test_session_key[i] = i + 100;
    }
    memcpy(ctx.session_key, test_session_key, SC_SESSION_KEY_SIZE);
    
    // 2. Test data preparation
    uint8_t plaintext[] = "Hello";
    size_t plaintext_len = sizeof(plaintext) - 1;
    uint8_t plaintext_with_crc[plaintext_len + SC_CRC32_SIZE];
    uint8_t combined_output[plaintext_len + SC_CRC32_SIZE + SC_TAG_SIZE];
    
    printf("Test data:\n");
    printf("  Plaintext: '%.*s' (%zu bytes)\n", (int)plaintext_len, plaintext, plaintext_len);
    printf("  Total buffer size: %zu bytes\n", plaintext_len + SC_CRC32_SIZE + SC_TAG_SIZE);
    
    // 3. Test CRC32 calculation
    memcpy(plaintext_with_crc, plaintext, plaintext_len);
    printf("\nStep 1: Copying plaintext... ✓\n");
    
    // Import CRC32 function
    extern uint32_t crc32_calc(const uint8_t *data, size_t len);
    uint32_t crc = crc32_calc(plaintext, plaintext_len);
    printf("Step 2: CRC32 calculation... 0x%08x\n", crc);
    
    plaintext_with_crc[plaintext_len] = (crc >> 0) & 0xFF;
    plaintext_with_crc[plaintext_len + 1] = (crc >> 8) & 0xFF;
    plaintext_with_crc[plaintext_len + 2] = (crc >> 16) & 0xFF;
    plaintext_with_crc[plaintext_len + 3] = (crc >> 24) & 0xFF;
    printf("Step 3: Adding CRC to data... ✓\n");
    
    // 4. Test AES key setup
    struct tc_aes_key_sched_struct sched;
    printf("\nStep 4: AES key setup...\n");
    printf("  Session key: ");
    for (int i = 0; i < SC_SESSION_KEY_SIZE; i++) {
        printf("%02x ", ctx.session_key[i]);
    }
    printf("\n");
    
    int aes_result = tc_aes128_set_encrypt_key(&sched, ctx.session_key);
    printf("  tc_aes128_set_encrypt_key returned: %d (should be 1)\n", aes_result);
    
    if (aes_result != 1) {
        printf("❌ AES key setup failed!\n");
        return 1;
    }
    printf("✓ AES key setup successful\n");
    
    // 5. Test nonce building
    uint8_t nonce[SC_NONCE_SIZE];
    printf("\nStep 5: Building nonce...\n");
    printf("  Counter: %llu\n", (unsigned long long)ctx.tx_counter);
    
    // Simple nonce building (mimicking sc_build_nonce)
    for (int i = 0; i < 4; i++) {
        nonce[i] = (ctx.tx_counter >> (i * 8)) & 0xFF;
    }
    for (int i = 4; i < SC_NONCE_SIZE; i++) {
        nonce[i] = i + 0x10;
    }
    printf("  Nonce: ");
    for (int i = 0; i < SC_NONCE_SIZE; i++) {
        printf("%02x ", nonce[i]);
    }
    printf("\n");
    
    // 6. Test CCM config
    struct tc_ccm_mode_struct ccm_state;
    printf("\nStep 6: CCM config...\n");
    printf("  Nonce size: %d, Tag size: %d\n", SC_NONCE_SIZE, SC_TAG_SIZE);
    
    int ccm_result = tc_ccm_config(&ccm_state, &sched, nonce, SC_NONCE_SIZE, SC_TAG_SIZE);
    printf("  tc_ccm_config returned: %d (should be 1)\n", ccm_result);
    
    if (ccm_result != 1) {
        printf("❌ CCM config failed!\n");
        return 1;
    }
    printf("✓ CCM config successful\n");
    
    // 7. Test encryption
    printf("\nStep 7: CCM encryption...\n");
    size_t total_plaintext_len = plaintext_len + SC_CRC32_SIZE;
    printf("  Total plaintext length: %zu bytes\n", total_plaintext_len);
    
    int enc_result = tc_ccm_generation_encryption(combined_output, sizeof(combined_output),
                                                  NULL, 0, /* no associated data */
                                                  plaintext_with_crc, total_plaintext_len,
                                                  &ccm_state);
    printf("  tc_ccm_generation_encryption returned: %d (should be 1)\n", enc_result);
    
    if (enc_result != 1) {
        printf("❌ CCM encryption failed!\n");
        return 1;
    }
    printf("🎉 CCM encryption successful!\n");
    printf("  Ciphertext length: %zu bytes\n", total_plaintext_len + SC_TAG_SIZE);
    
    // 8. Test decryption
    printf("\nStep 8: CCM decryption...\n");
    uint8_t decrypted[total_plaintext_len];
    
    int dec_result = tc_ccm_decryption_verification(decrypted, total_plaintext_len,
                                                   NULL, 0, combined_output, total_plaintext_len + SC_TAG_SIZE,
                                                   &ccm_state);
    printf("  tc_ccm_decryption_verification returned: %d (should be 1)\n", dec_result);
    
    if (dec_result != 1) {
        printf("❌ CCM decryption failed!\n");
        return 1;
    }
    printf("🎉 CCM decryption successful!\n");
    
    // Verify result
    printf("\nVerification:\n");
    printf("  Original: '%.*s'\n", (int)plaintext_len, plaintext);
    printf("  Decrypted: '%.*s'\n", (int)plaintext_len, decrypted);
    
    if (memcmp(plaintext, decrypted, plaintext_len) == 0) {
        printf("✓ Decrypted data matches original!\n");
    } else {
        printf("❌ Decrypted data doesn't match original!\n");
    }
    
    printf("\n=== All crypto operations successful! ===\n");
    return 0;
}