// test_stream_sign.c — tests for streaming Ed25519 signature
#include "../src/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, "=== Streaming Sign Test ===");

    struct SC_MYKEYS keys_a;
    CHECK(sc_generate_keypair(&keys_a) == SC_OK, "generate keypair A");
    
    sc_context_t ctx_a;
    CHECK(sc_init_ctx(&ctx_a, &keys_a) == SC_OK, "init ctx A");

    uint8_t ed25519_pubkey_a[SC_PUBKEY_SIZE];
    CHECK(sc_derive_ed25519_pubkey(keys_a.private_key, ed25519_pubkey_a) == SC_OK, "derive ed25519 pubkey from x25519");

    // ==========================================
    // Test 1: basic sign/verify — single chunk
    // ==========================================
    {
        uint8_t data[1024];
        for (int i = 0; i < (int)sizeof(data); i++) data[i] = (uint8_t)(i * 7 + 31);

        struct sc_stream_sign_state signer, verifier;
        CHECK(sc_stream_sign_init(&ctx_a, &signer) == SC_OK, "sign_init");
        CHECK(sc_stream_sign_verify_init(&verifier, ed25519_pubkey_a) == SC_OK, "verify_init");

        CHECK(sc_stream_sign_update(&signer, data, sizeof(data)) == SC_OK, "sign_update 1024B");
        CHECK(sc_stream_sign_update(&verifier, data, sizeof(data)) == SC_OK, "verify_update 1024B");

        uint8_t sig[SC_SIGN_SIZE];
        size_t sig_len = sizeof(sig);
        CHECK(sc_stream_sign_final(&signer, sig, &sig_len) == SC_OK, "sign_final");
        CHECK(sig_len == SC_SIGN_SIZE, "signature size 64");

        CHECK(sc_stream_sign_verify(&verifier, sig, sig_len) == SC_OK, "verify OK");
    }

    // ==========================================
    // Test 2: incremental update — many small chunks
    // ==========================================
    {
        uint8_t data[2048];
        for (int i = 0; i < (int)sizeof(data); i++) data[i] = (uint8_t)(i * 3 + 17);

        struct sc_stream_sign_state signer, verifier;
        CHECK(sc_stream_sign_init(&ctx_a, &signer) == SC_OK, "inc: sign_init");
        CHECK(sc_stream_sign_verify_init(&verifier, ed25519_pubkey_a) == SC_OK, "inc: verify_init");

        for (int i = 0; i < (int)sizeof(data); i += 7) {
            int chunk = ((int)sizeof(data) - i < 7) ? ((int)sizeof(data) - i) : 7;
            CHECK(sc_stream_sign_update(&signer, data + i, chunk) == SC_OK, "inc: sign_update");
            CHECK(sc_stream_sign_update(&verifier, data + i, chunk) == SC_OK, "inc: verify_update");
        }

        uint8_t sig[SC_SIGN_SIZE];
        size_t sig_len = sizeof(sig);
        CHECK(sc_stream_sign_final(&signer, sig, &sig_len) == SC_OK, "inc: sign_final");

        CHECK(sc_stream_sign_verify(&verifier, sig, sig_len) == SC_OK, "inc: verify OK");
    }

    // ==========================================
    // Test 3: tampered data — verify MUST fail
    // ==========================================
    {
        uint8_t data[512];
        for (int i = 0; i < (int)sizeof(data); i++) data[i] = (uint8_t)(i * 5 + 7);

        struct sc_stream_sign_state signer, verifier;
        CHECK(sc_stream_sign_init(&ctx_a, &signer) == SC_OK, "tamper: sign_init");
        CHECK(sc_stream_sign_verify_init(&verifier, ed25519_pubkey_a) == SC_OK, "tamper: verify_init");

        CHECK(sc_stream_sign_update(&signer, data, sizeof(data)) == SC_OK, "tamper: sign_update");

        uint8_t sig[SC_SIGN_SIZE];
        size_t sig_len = sizeof(sig);
        CHECK(sc_stream_sign_final(&signer, sig, &sig_len) == SC_OK, "tamper: sign_final");

        data[0] ^= 0x01;
        CHECK(sc_stream_sign_update(&verifier, data, sizeof(data)) == SC_OK, "tamper: verify_update (modified)");
        CHECK(sc_stream_sign_verify(&verifier, sig, sig_len) == SC_ERR_AUTH_FAILED, "tamper: verify FAIL (expected)");
    }

    // ==========================================
    // Test 4: wrong public key — verify MUST fail
    // ==========================================
    {
        struct SC_MYKEYS keys_b;
        CHECK(sc_generate_keypair(&keys_b) == SC_OK, "generate keypair B");
        uint8_t pubkey_b[SC_PUBKEY_SIZE];
        CHECK(sc_derive_ed25519_pubkey(keys_b.private_key, pubkey_b) == SC_OK, "derive pubkey B");

        uint8_t data[256];
        for (int i = 0; i < (int)sizeof(data); i++) data[i] = (uint8_t)i;

        struct sc_stream_sign_state signer, verifier;
        CHECK(sc_stream_sign_init(&ctx_a, &signer) == SC_OK, "wrongkey: sign_init (A)");
        CHECK(sc_stream_sign_verify_init(&verifier, pubkey_b) == SC_OK, "wrongkey: verify_init (B's key)");

        CHECK(sc_stream_sign_update(&signer, data, sizeof(data)) == SC_OK, "wrongkey: sign_update");
        CHECK(sc_stream_sign_update(&verifier, data, sizeof(data)) == SC_OK, "wrongkey: verify_update");

        uint8_t sig[SC_SIGN_SIZE];
        size_t sig_len = sizeof(sig);
        CHECK(sc_stream_sign_final(&signer, sig, &sig_len) == SC_OK, "wrongkey: sign_final");

        CHECK(sc_stream_sign_verify(&verifier, sig, sig_len) == SC_ERR_AUTH_FAILED, "wrongkey: verify FAIL (expected)");
    }

    // ==========================================
    // Test 5: empty data — sign/verify empty message
    // ==========================================
    {
        struct sc_stream_sign_state signer, verifier;
        CHECK(sc_stream_sign_init(&ctx_a, &signer) == SC_OK, "empty: sign_init");
        CHECK(sc_stream_sign_verify_init(&verifier, ed25519_pubkey_a) == SC_OK, "empty: verify_init");

        uint8_t sig[SC_SIGN_SIZE];
        size_t sig_len = sizeof(sig);
        CHECK(sc_stream_sign_final(&signer, sig, &sig_len) == SC_OK, "empty: sign_final (no data)");

        CHECK(sc_stream_sign_verify(&verifier, sig, sig_len) == SC_OK, "empty: verify OK");
    }

    // ==========================================
    // Test 6: update with len=0 — should be no-op
    // ==========================================
    {
        uint8_t data[64];
        for (int i = 0; i < (int)sizeof(data); i++) data[i] = (uint8_t)(i * 11);

        struct sc_stream_sign_state signer;
        CHECK(sc_stream_sign_init(&ctx_a, &signer) == SC_OK, "zerolen: sign_init");
        CHECK(sc_stream_sign_update(&signer, NULL, 0) == SC_OK, "zerolen: NULL+0 OK");
        CHECK(sc_stream_sign_update(&signer, data, 0) == SC_OK, "zerolen: data+0 OK");
        CHECK(sc_stream_sign_update(&signer, data, sizeof(data)) == SC_OK, "zerolen: sign_update 64B");

        uint8_t sig[SC_SIGN_SIZE];
        size_t sig_len = sizeof(sig);
        CHECK(sc_stream_sign_final(&signer, sig, &sig_len) == SC_OK, "zerolen: sign_final");
    }

    // ==========================================
    // Test 7: error handling — init
    // ==========================================
    {
        struct sc_stream_sign_state state;
        memset(&state, 0, sizeof(state));
        CHECK(sc_stream_sign_init(NULL, &state) == SC_ERR_INVALID_ARG, "err: init NULL ctx");
        CHECK(sc_stream_sign_init(&ctx_a, NULL) == SC_ERR_INVALID_ARG, "err: init NULL state");
        
        sc_context_t no_pk_ctx;
        memset(&no_pk_ctx, 0, sizeof(no_pk_ctx));
        CHECK(sc_stream_sign_init(&no_pk_ctx, &state) == SC_ERR_NOT_INITIALIZED, "err: init uninitialized ctx");

        CHECK(sc_stream_sign_verify_init(NULL, ed25519_pubkey_a) == SC_ERR_INVALID_ARG, "err: verify_init NULL state");
        CHECK(sc_stream_sign_verify_init(&state, NULL) == SC_ERR_INVALID_ARG, "err: verify_init NULL pubkey");
    }

    // ==========================================
    // Test 8: error handling — update
    // ==========================================
    {
        struct sc_stream_sign_state state;
        memset(&state, 0, sizeof(state));
        CHECK(sc_stream_sign_update(NULL, (uint8_t*)"x", 1) == SC_ERR_INVALID_ARG, "err: update NULL state");
        CHECK(sc_stream_sign_update(&state, NULL, 1) == SC_ERR_INVALID_ARG, "err: update NULL data");
        CHECK(sc_stream_sign_update(&state, (uint8_t*)"x", 1) == SC_ERR_NOT_INITIALIZED, "err: update not init");
    }

    // ==========================================
    // Test 9: error handling — final/verify on wrong init
    // ==========================================
    {
        struct sc_stream_sign_state verifier;
        CHECK(sc_stream_sign_verify_init(&verifier, ed25519_pubkey_a) == SC_OK, "err: verify_init for final test");

        uint8_t sig[SC_SIGN_SIZE] = {0};
        size_t sig_len = sizeof(sig);
        CHECK(sc_stream_sign_final(&verifier, sig, &sig_len) == SC_ERR_NOT_INITIALIZED, "err: final on verify state");
        sc_stream_sign_cleanup(&verifier);
    }
    {
        struct sc_stream_sign_state signer;
        CHECK(sc_stream_sign_init(&ctx_a, &signer) == SC_OK, "err: sign_init for verify test");
        uint8_t sig[SC_SIGN_SIZE] = {0};
        CHECK(sc_stream_sign_verify(&signer, sig, SC_SIGN_SIZE) == SC_ERR_NOT_INITIALIZED, "err: verify on sign state");
        sc_stream_sign_cleanup(&signer);
    }

    // ==========================================
    // Test 10: cleanup
    // ==========================================
    {
        struct sc_stream_sign_state state;
        memset(&state, 0, sizeof(state));
        sc_stream_sign_cleanup(&state);    // uninitialized — no crash
        sc_stream_sign_cleanup(NULL);      // NULL — no crash

        CHECK(sc_stream_sign_init(&ctx_a, &state) == SC_OK, "cleanup: init");
        sc_stream_sign_cleanup(&state);
        CHECK(state.initialized == 0, "cleanup: initialized cleared");
    }

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