#include "naut/hash.h"
#include "naut/rc4.h"
#include "naut/merkle.h"
#include "test.h"
#include <string.h>
#include <stdlib.h>

static int hexeq(const uint8_t *got, const char *hex, size_t n) {
    for (size_t i = 0; i < n; i++) {
        unsigned b;
        sscanf(hex + i*2, "%2x", &b);
        if (got[i] != (uint8_t)b) return 0;
    }
    return 1;
}

int main(void) {
    uint8_t d[32];

    /* --- SHA-1 (FIPS) --- */
    naut_sha1("abc", 3, d);
    CHECK(hexeq(d, "a9993e364706816aba3e25717850c26c9cd0d89d", 20));
    naut_sha1("", 0, d);
    CHECK(hexeq(d, "da39a3ee5e6b4b0d3255bfef95601890afd80709", 20));

    /* --- SHA-256 (FIPS) --- */
    fprintf(stderr, "sha256 backend: %s\n", naut_sha256_backend());
    naut_sha256("abc", 3, d);
    CHECK(hexeq(d, "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad", 32));
    naut_sha256("", 0, d);
    CHECK(hexeq(d, "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", 32));

    /* multi-block + streaming: 1,000,000 'a' fed in odd-sized chunks */
    {
        naut_sha256_ctx c; naut_sha256_init(&c);
        char chunk[7]; memset(chunk, 'a', sizeof chunk);
        size_t left = 1000000;
        while (left) { size_t n = left < sizeof chunk ? left : sizeof chunk;
                       naut_sha256_update(&c, chunk, n); left -= n; }
        naut_sha256_final(&c, d);
        CHECK(hexeq(d, "cdc76e5c9914fb9281a1c7e284d73e67f1809a48a497200e046d39ccc7112cd0", 32));
    }

    /* --- RC4 (Wikipedia test vectors, drop=0) --- */
    {
        naut_rc4 c; uint8_t buf[16];
        memcpy(buf, "Plaintext", 9);
        naut_rc4_init(&c, "Key", 3, 0); naut_rc4_xor(&c, buf, 9);
        CHECK(hexeq(buf, "bbf316e8d940af0ad3", 9));

        memcpy(buf, "pedia", 5);
        naut_rc4_init(&c, "Wiki", 4, 0); naut_rc4_xor(&c, buf, 5);
        CHECK(hexeq(buf, "1021bf0420", 5));

        /* round-trip with MSE-style drop=1024 */
        uint8_t msg[64], orig[64];
        for (int i = 0; i < 64; i++) orig[i] = msg[i] = (uint8_t)(i * 7);
        naut_rc4_init(&c, "secretkey", 9, 1024); naut_rc4_xor(&c, msg, 64);
        CHECK(memcmp(msg, orig, 64) != 0);              /* actually encrypted */
        naut_rc4_init(&c, "secretkey", 9, 1024); naut_rc4_xor(&c, msg, 64);
        CHECK(memcmp(msg, orig, 64) == 0);              /* decrypts back */
    }

    /* --- Merkle (BEP-52 structure) --- */
    {
        /* 3 leaf hashes -> pad to 4 with a zero leaf */
        uint8_t leaves[3*32];
        for (int i = 0; i < 3; i++) { char s[2] = { (char)('a'+i), 0 };
            naut_sha256(s, 1, leaves + i*32); }

        uint8_t root[32], expect[32];
        CHECK(naut_merkle_root(leaves, 3, root) == NAUT_OK);

        /* recompute by hand: n0=H(l0||l1), n1=H(l2||zero), root=H(n0||n1) */
        uint8_t zero[32] = {0}, n0[32], n1[32], cat[64];
        memcpy(cat, leaves+0, 32); memcpy(cat+32, leaves+32, 32); naut_sha256(cat, 64, n0);
        memcpy(cat, leaves+64, 32); memcpy(cat+32, zero, 32);     naut_sha256(cat, 64, n1);
        memcpy(cat, n0, 32); memcpy(cat+32, n1, 32);              naut_sha256(cat, 64, expect);
        CHECK(memcmp(root, expect, 32) == 0);

        /* single leaf: root == leaf */
        CHECK(naut_merkle_root(leaves, 1, root) == NAUT_OK);
        CHECK(memcmp(root, leaves, 32) == 0);

        /* leaves from a 40 KiB buffer -> 3 leaves (16K,16K,8K) */
        size_t len = 40*1024;
        uint8_t *blob = malloc(len); memset(blob, 0x5a, len);
        uint8_t lv[3*32];
        CHECK_EQ(naut_merkle_leaves(blob, len, lv), 3);
        free(blob);
    }

    TEST_MAIN_END();
}