#include <openssl/sha.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "../sha256/sha256_backend.h"
#include "../utils.h"

static uint32_t xorshift32(uint32_t *s) {
    uint32_t x = *s;
    x ^= x << 13;
    x ^= x >> 17;
    x ^= x << 5;
    *s = x;
    return x;
}

static inline void write_u32_le(uint8_t *dst, uint32_t v) {
    dst[0] = (uint8_t)(v & 0xFFU);
    dst[1] = (uint8_t)((v >> 8) & 0xFFU);
    dst[2] = (uint8_t)((v >> 16) & 0xFFU);
    dst[3] = (uint8_t)((v >> 24) & 0xFFU);
}

static void fill_random(uint8_t *dst, size_t n, uint32_t *seed) {
    size_t i;
    for (i = 0; i < n; i++) {
        dst[i] = (uint8_t)(xorshift32(seed) & 0xFFU);
    }
}

static void backend_double_sha256_80(const uint8_t header80[80], uint8_t out[32]) {
    sha256_state_t init_state;
    sha256_state_t st1;
    sha256_state_t st2;
    uint8_t block1[64];
    uint8_t block2[64];

    sha256_state_init(&init_state);

    st1 = init_state;
    sha256_transform_fast(&st1, header80);

    memset(block1, 0, sizeof(block1));
    memcpy(block1, header80 + 64, 16);
    block1[16] = 0x80;
    block1[62] = 0x02;
    block1[63] = 0x80;
    sha256_transform_fast(&st1, block1);

    memset(block2, 0, sizeof(block2));
    sha256_state_to_digest(&st1, block2);
    block2[32] = 0x80;
    block2[62] = 0x01;
    block2[63] = 0x00;

    st2 = init_state;
    sha256_transform_fast(&st2, block2);
    sha256_state_to_digest(&st2, out);
}

static int test_transform_equivalence_random(void) {
    uint32_t seed = 0x12345678U;
    int i;

    for (i = 0; i < 10000; i++) {
        uint8_t block[64];
        SHA256_CTX ref;
        sha256_state_t test;

        fill_random(block, sizeof(block), &seed);

        SHA256_Init(&ref);
        sha256_state_init(&test);
        SHA256_Transform(&ref, block);
        sha256_transform_fast(&test, block);

        if (memcmp(ref.h, test.h, sizeof(ref.h)) != 0) {
            fprintf(stderr, "[test_sha256_backend] transform mismatch at iter=%d\n", i);
            return 0;
        }
    }

    return 1;
}

static int test_transform_2way_equivalence(void) {
    uint32_t seed = 0xCAFEBABEU;
    int i;

    for (i = 0; i < 5000; i++) {
        uint8_t block_a[64];
        uint8_t block_b[64];
        SHA256_CTX ref_a;
        SHA256_CTX ref_b;
        sha256_state_t test_a;
        sha256_state_t test_b;

        fill_random(block_a, sizeof(block_a), &seed);
        fill_random(block_b, sizeof(block_b), &seed);

        SHA256_Init(&ref_a);
        SHA256_Init(&ref_b);
        sha256_state_init(&test_a);
        sha256_state_init(&test_b);

        SHA256_Transform(&ref_a, block_a);
        SHA256_Transform(&ref_b, block_b);
        sha256_transform_fast_2way(&test_a, block_a, &test_b, block_b);

        if (memcmp(ref_a.h, test_a.h, sizeof(ref_a.h)) != 0 ||
            memcmp(ref_b.h, test_b.h, sizeof(ref_b.h)) != 0) {
            fprintf(stderr, "[test_sha256_backend] 2way mismatch at iter=%d\n", i);
            return 0;
        }
    }

    return 1;
}

static int test_double_sha256_80_equivalence(void) {
    uint32_t seed = 0xA55AA55AU;
    int i;

    for (i = 0; i < 2048; i++) {
        uint8_t header80[80];
        uint8_t ref[32];
        uint8_t test[32];

        fill_random(header80, sizeof(header80), &seed);
        double_sha256(header80, sizeof(header80), ref);
        backend_double_sha256_80(header80, test);

        if (memcmp(ref, test, sizeof(ref)) != 0) {
            fprintf(stderr, "[test_sha256_backend] sha256d80 mismatch at iter=%d\n", i);
            return 0;
        }
    }

    return 1;
}

static int test_sha256d80_4way_100k_nonces(void) {
    uint8_t header76[76];
    sha256d80_midstate_t mid;
    uint32_t base = 0x12340000U;
    int i;

    for (i = 0; i < 76; i++) {
        header76[i] = (uint8_t)((i * 19 + 5) & 0xFF);
    }

    sha256d80_midstate_init(&mid, header76);

    for (i = 0; i < 25000; i++) {
        uint32_t start_nonce = base + (uint32_t)(i * 4);
        sha256_state_t states[4];
        int lane;

        sha256d80_hash_4way(&mid, start_nonce, states);

        for (lane = 0; lane < 4; lane++) {
            uint8_t header80[80];
            uint8_t ref[32];
            uint8_t got[32];
            uint32_t nonce = start_nonce + (uint32_t)lane;

            memcpy(header80, header76, 76);
            write_u32_le(header80 + 76, nonce);
            double_sha256(header80, sizeof(header80), ref);
            sha256_state_to_digest(&states[lane], got);

            if (memcmp(ref, got, 32) != 0) {
                fprintf(stderr, "[test_sha256_backend] 4way nonce mismatch nonce=%u lane=%d\n", nonce, lane);
                return 0;
            }
        }
    }

    return 1;
}

static int test_sha256d80_scan_hitmask_basic(void) {
    uint8_t header76[76];
    sha256d80_midstate_t mid;
    sha256_state_t states[4];
    uint32_t all_max[8];
    uint32_t all_zero[8];
    uint32_t mask_max;
    uint32_t mask_zero;
    int i;

    for (i = 0; i < 76; i++) {
        header76[i] = (uint8_t)((i * 11 + 1) & 0xFF);
    }
    for (i = 0; i < 8; i++) {
        all_max[i] = 0xFFFFFFFFU;
        all_zero[i] = 0x00000000U;
    }

    sha256d80_midstate_init(&mid, header76);
    mask_max = sha256d80_scan_4way(&mid, 0xABC00000U, all_max, states);
    mask_zero = sha256d80_scan_4way(&mid, 0xABC00000U, all_zero, states);

    return mask_max == 0xFU && mask_zero == 0U;
}

int main(void) {
    if (!test_transform_equivalence_random()) {
        return 1;
    }
    if (!test_transform_2way_equivalence()) {
        return 1;
    }
    if (!test_double_sha256_80_equivalence()) {
        return 1;
    }
    if (!test_sha256d80_4way_100k_nonces()) {
        return 1;
    }
    if (!test_sha256d80_scan_hitmask_basic()) {
        return 1;
    }

    printf("test_sha256_backend: OK\n");
    return 0;
}