/* test_sc_lib.c - Test for Secure Channel library */

#include "sc_lib.h"
#include <tinycrypt/ecc.h>
#include <tinycrypt/ecc_platform_specific.h>
#include <stdio.h>
#include <string.h>
#include <stdint.h>

#define TEST_ASSERT(cond, msg) \
    do { \
        if (!(cond)) { \
            printf("FAIL: %s (line %d)\n", msg, __LINE__); \
            return 1; \
        } else { \
            printf("PASS: %s\n", msg); \
        } \
    } while(0)

static int test_key_generation(void)
{
    sc_context_t ctx;
    sc_status_t status;
    
    memset(&ctx, 0, sizeof(ctx));
    
    status = sc_generate_keypair(&ctx);
    TEST_ASSERT(status == SC_OK, "sc_generate_keypair");
    TEST_ASSERT(ctx.initialized == 1, "ctx.initialized set");
    TEST_ASSERT(ctx.peer_key_set == 0, "peer_key_set initially 0");
    TEST_ASSERT(ctx.session_ready == 0, "session_ready initially 0");
    
    return 0;
}

static int test_key_exchange(void)
{
    sc_context_t client_ctx, server_ctx;
    sc_status_t status;
    
    memset(&client_ctx, 0, sizeof(client_ctx));
    memset(&server_ctx, 0, sizeof(server_ctx));
    
    /* Generate key pairs */
    status = sc_generate_keypair(&client_ctx);
    TEST_ASSERT(status == SC_OK, "client key generation");
    
    status = sc_generate_keypair(&server_ctx);
    TEST_ASSERT(status == SC_OK, "server key generation");
    
    /* Exchange public keys */
    status = sc_set_peer_public_key(&client_ctx, server_ctx.public_key);
    TEST_ASSERT(status == SC_OK, "client set peer key");
    TEST_ASSERT(client_ctx.peer_key_set == 1, "client peer_key_set");
    TEST_ASSERT(client_ctx.session_ready == 1, "client session_ready");
    
    status = sc_set_peer_public_key(&server_ctx, client_ctx.public_key);
    TEST_ASSERT(status == SC_OK, "server set peer key");
    TEST_ASSERT(server_ctx.peer_key_set == 1, "server peer_key_set");
    TEST_ASSERT(server_ctx.session_ready == 1, "server session_ready");
    
    /* Session keys should be identical */
    TEST_ASSERT(memcmp(client_ctx.session_key, server_ctx.session_key,
                       SC_SESSION_KEY_SIZE) == 0,
                "session keys match");
    
    return 0;
}

static int test_encrypt_decrypt(void)
{
    sc_context_t client_ctx, server_ctx;
    sc_status_t status;
    const char *plaintext = "Hello, secure channel!";
    size_t plaintext_len = strlen(plaintext) + 1;
    uint8_t ciphertext[256];
    size_t ciphertext_len;
    uint8_t decrypted[256];
    size_t decrypted_len;
    
    memset(&client_ctx, 0, sizeof(client_ctx));
    memset(&server_ctx, 0, sizeof(server_ctx));
    
    /* Setup secure channel */
    status = sc_generate_keypair(&client_ctx);
    TEST_ASSERT(status == SC_OK, "client key generation");
    
    status = sc_generate_keypair(&server_ctx);
    TEST_ASSERT(status == SC_OK, "server key generation");
    
    status = sc_set_peer_public_key(&client_ctx, server_ctx.public_key);
    TEST_ASSERT(status == SC_OK, "client set peer key");
    
    status = sc_set_peer_public_key(&server_ctx, client_ctx.public_key);
    TEST_ASSERT(status == SC_OK, "server set peer key");
    
    /* Client encrypts */
    status = sc_encrypt(&client_ctx, (const uint8_t *)plaintext, plaintext_len,
                        ciphertext, &ciphertext_len);
    TEST_ASSERT(status == SC_OK, "encryption");
    TEST_ASSERT(ciphertext_len == plaintext_len + SC_MAX_OVERHEAD, "ciphertext length includes overhead");
    
    /* Server decrypts */
    status = sc_decrypt(&server_ctx, ciphertext, ciphertext_len, decrypted, &decrypted_len);
    TEST_ASSERT(status == SC_OK, "decryption");
    TEST_ASSERT(decrypted_len == plaintext_len, "decrypted length matches original");
    
    /* Verify decrypted matches original */
    TEST_ASSERT(memcmp(plaintext, decrypted, plaintext_len) == 0,
                "decrypted matches original");
    
    /* Verify tag verification fails with modified ciphertext */
    uint8_t corrupted_ciphertext[256];
    memcpy(corrupted_ciphertext, ciphertext, ciphertext_len);
    corrupted_ciphertext[10] ^= 0xAA;  /* Изменяем байт шифротекста */
    status = sc_decrypt(&server_ctx, corrupted_ciphertext, ciphertext_len, decrypted, &decrypted_len);
    TEST_ASSERT(status == SC_ERR_AUTH_FAILED, "modified ciphertext detection");
    
    /* Verify counter increment */
    TEST_ASSERT(client_ctx.tx_counter == 1, "client tx_counter incremented");
    TEST_ASSERT(server_ctx.rx_counter == 1, "server rx_counter incremented");
    
    return 0;
}

static int test_error_handling(void)
{
    sc_context_t ctx;
    uint8_t buffer[32];
    size_t out_len;
    sc_status_t status;
    
    memset(&ctx, 0, sizeof(ctx));
    
    /* Test NULL context */
    status = sc_generate_keypair(NULL);
    TEST_ASSERT(status == SC_ERR_INVALID_ARG, "NULL context detection");
    
    /* Test uninitialized context */
    memset(&ctx, 0, sizeof(ctx));
    status = sc_set_peer_public_key(&ctx, buffer);
    TEST_ASSERT(status == SC_ERR_NOT_INITIALIZED, "uninitialized context");
    
    /* Test encryption without session */
    memset(&ctx, 0, sizeof(ctx));
    status = sc_generate_keypair(&ctx);
    TEST_ASSERT(status == SC_OK, "key generation for error test");
    
    status = sc_encrypt(&ctx, buffer, sizeof(buffer), buffer, &out_len);
    TEST_ASSERT(status == SC_ERR_NOT_INITIALIZED, "encrypt without session");
    
    return 0;
}

int main(void)
{
    int result = 0;
    
    printf("=== Secure Channel Library Test ===\n");
    
    /* Initialize TinyCrypt RNG */
    uECC_set_rng(&default_CSPRNG);
    
    result |= test_key_generation();
    result |= test_key_exchange();
    result |= test_encrypt_decrypt();
    result |= test_error_handling();
    
    if (result == 0) {
        printf("\n=== All tests PASSED ===\n");
    } else {
        printf("\n=== Some tests FAILED ===\n");
    }
    
    return result;
}