// test_stream_cipher.c — tests for streaming cipher (AES-128-CTR)
#include "secure_channel.h"
#include "../lib/debug_config.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

static int test_failed = 0;

#define CHECK(expr, msg) do { \
    if (!(expr)) { \
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "FAIL: %s", msg); \
        test_failed = 1; \
    } else { \
        DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "PASS: %s", msg); \
    } \
} while(0)

int main(void) {
    debug_config_init();
    debug_set_level(DEBUG_LEVEL_INFO);
    debug_set_categories(DEBUG_CATEGORY_CRYPTO);

    DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "=== Stream Cipher Test ===");

    struct SC_MYKEYS keys_a, keys_b;
    CHECK(sc_generate_keypair(&keys_a) == SC_OK, "generate keypair A");
    CHECK(sc_generate_keypair(&keys_b) == SC_OK, "generate keypair B");

    sc_context_t ctx_a, ctx_b;
    CHECK(sc_init_ctx(&ctx_a, &keys_a) == SC_OK, "init ctx A");
    CHECK(sc_init_ctx(&ctx_b, &keys_b) == SC_OK, "init ctx B");

    CHECK(sc_set_peer_public_key(&ctx_a, keys_b.public_key, SC_PEER_PUBKEY_BIN) == SC_OK, "set peer key A<-B");
    CHECK(sc_set_peer_public_key(&ctx_b, keys_a.public_key, SC_PEER_PUBKEY_BIN) == SC_OK, "set peer key B<-A");

    struct sc_stream_state send_a, recv_a, send_b, recv_b;
    CHECK(sc_stream_init(&ctx_a, &send_a, 0) == SC_OK, "init send_A (id=0)");
    CHECK(sc_stream_init(&ctx_a, &recv_a, 1) == SC_OK, "init recv_A (id=1)");
    CHECK(sc_stream_init(&ctx_b, &recv_b, 0) == SC_OK, "init recv_B (id=0)");
    CHECK(sc_stream_init(&ctx_b, &send_b, 1) == SC_OK, "init send_B (id=1)");

    uint8_t original[2048], data[2048];
    for (int i = 0; i < (int)sizeof(original); i++) original[i] = (uint8_t)(i * 3 + 17);

    // Test 1: A->B, single chunk, 1024 bytes
    memcpy(data, original, 1024);
    CHECK(sc_stream_xor(&send_a, data, 1024) == SC_OK, "encrypt 1024B A->B");
    CHECK(sc_stream_xor(&recv_b, data, 1024) == SC_OK, "decrypt 1024B A->B");
    CHECK(memcmp(data, original, 1024) == 0, "round-trip A->B matches");

    // Test 2: B->A, single chunk
    memcpy(data, original, 1024);
    CHECK(sc_stream_xor(&send_b, data, 1024) == SC_OK, "encrypt 1024B B->A");
    CHECK(sc_stream_xor(&recv_a, data, 1024) == SC_OK, "decrypt 1024B B->A");
    CHECK(memcmp(data, original, 1024) == 0, "round-trip B->A matches");

    // Test 3: incremental feeding — encrypt in 4 chunks, decrypt in 2
    memcpy(data, original, 1024);
    CHECK(sc_stream_xor(&send_a, data, 7) == SC_OK, "inc_A enc 7B");
    CHECK(sc_stream_xor(&send_a, data + 7, 33) == SC_OK, "inc_A enc 33B");
    CHECK(sc_stream_xor(&send_a, data + 40, 1) == SC_OK, "inc_A enc 1B");
    CHECK(sc_stream_xor(&send_a, data + 41, 1024 - 41) == SC_OK, "inc_A enc 983B");
    CHECK(sc_stream_xor(&recv_b, data, 513) == SC_OK, "inc_B dec 513B");
    CHECK(sc_stream_xor(&recv_b, data + 513, 1024 - 513) == SC_OK, "inc_B dec 511B");
    CHECK(memcmp(data, original, 1024) == 0, "inc round-trip matches");

    // Test 4: all sizes 1..255 — catches alignment/keystream issues
    for (size_t sz = 1; sz <= 255; sz++) {
        uint8_t buf[256], cpy[256];
        for (size_t i = 0; i < sz; i++) buf[i] = (uint8_t)(sz ^ i);
        memcpy(cpy, buf, sz);
        CHECK(sc_stream_xor(&send_a, buf, sz) == SC_OK, "various-size enc");
        CHECK(sc_stream_xor(&recv_b, buf, sz) == SC_OK, "various-size dec");
        CHECK(memcmp(buf, cpy, sz) == 0, "various-size match");
    }

    // Test 5: 2048 bytes (= 128 AES blocks, whole number)
    memcpy(data, original, 2048);
    CHECK(sc_stream_xor(&send_a, data, 2048) == SC_OK, "enc 2048B");
    CHECK(sc_stream_xor(&recv_b, data, 2048) == SC_OK, "dec 2048B");
    CHECK(memcmp(data, original, 2048) == 0, "round-trip 2048 matches");

    // Test 6: A and B streams must produce DIFFERENT keystreams
    {
        uint8_t test_data[16];
        memset(test_data, 0xAA, 16);
        uint8_t a_out[16], b_out[16];
        memcpy(a_out, test_data, 16);
        memcpy(b_out, test_data, 16);
        CHECK(sc_stream_xor(&send_a, a_out, 16) == SC_OK, "xor on send_A");
        CHECK(sc_stream_xor(&send_b, b_out, 16) == SC_OK, "xor on send_B");
        CHECK(memcmp(a_out, b_out, 16) != 0, "stream_id=0 != stream_id=1 keystream");
        // Decrypt back
        CHECK(sc_stream_xor(&recv_b, a_out, 16) == SC_OK, "recv_B decrypt send_A data");
        CHECK(sc_stream_xor(&recv_a, b_out, 16) == SC_OK, "recv_A decrypt send_B data");
        CHECK(memcmp(a_out, test_data, 16) == 0, "recv_B restores original");
        CHECK(memcmp(b_out, test_data, 16) == 0, "recv_A restores original");
    }

    // Test 7: encrypt 1 byte — extreme alignment edge
    for (int j = 0; j < 50; j++) {
        uint8_t one_byte = (uint8_t)j;
        CHECK(sc_stream_xor(&send_a, &one_byte, 1) == SC_OK, "enc 1B");
        CHECK(sc_stream_xor(&recv_b, &one_byte, 1) == SC_OK, "dec 1B");
        CHECK(one_byte == (uint8_t)j, "1-byte round-trip");
    }

    // Test 8: error handling
    CHECK(sc_stream_xor(&send_a, NULL, 0) == SC_OK, "xor len=0 OK");
    CHECK(sc_stream_xor(&send_a, data, 0) == SC_OK, "xor data_len=0 OK");
    CHECK(sc_stream_init(&ctx_a, NULL, 0) == SC_ERR_INVALID_ARG, "init NULL state");
    CHECK(sc_stream_init(NULL, &send_a, 0) == SC_ERR_INVALID_ARG, "init NULL ctx");

    struct sc_stream_state uninit;
    memset(&uninit, 0, sizeof(uninit));
    CHECK(sc_stream_xor(&uninit, data, 1) == SC_ERR_NOT_INITIALIZED, "xor uninitialized");

    // У ctx без session_ready
    sc_context_t no_session_ctx;
    struct SC_MYKEYS dummy_keys;
    memset(&dummy_keys, 0, sizeof(dummy_keys));
    sc_init_ctx(&no_session_ctx, &dummy_keys);
    struct sc_stream_state bad_stream;
    CHECK(sc_stream_init(&no_session_ctx, &bad_stream, 0) == SC_ERR_NOT_INITIALIZED, "init without session");

    // Cleanup
    sc_stream_cleanup(&send_a);
    sc_stream_cleanup(&recv_a);
    sc_stream_cleanup(&send_b);
    sc_stream_cleanup(&recv_b);
    sc_stream_cleanup(NULL);
    sc_stream_cleanup(&uninit);

    if (test_failed) {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "=== Stream Cipher Test: FAILED ===");
        return 1;
    }
    DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "=== Stream Cipher Test: PASSED ===");
    return 0;
}