#include "unity/unity.h"
#include "zlib.h"
#include <stdint.h>
#include <stdlib.h>

/* Wrapper for local function provided in module */
extern z_crc_t test_x2nmodp(z_off64_t n, unsigned k);

void setUp(void) {
    /* Ensure CRC tables are initialized in dynamic-table builds */
    const z_crc_t *tbl = get_crc_table();
    TEST_ASSERT_NOT_NULL(tbl);
}

void tearDown(void) {
    /* nothing */
}

/* Helper to multiply two operator polynomials using public API:
   mult(opA, opB) = crc32_combine_op(opB, 0, opA) */
static uint32_t op_mul(uint32_t opA, uint32_t opB) {
    uLong prod = crc32_combine_op((uLong)opB, (uLong)0, (uLong)opA);
    return (uint32_t)prod;
}

void test_x2nmodp_identity_and_small_powers(void) {
    /* Identity: n=0 => x^0 == 0x80000000 irrespective of k */
    uint32_t id0 = (uint32_t)test_x2nmodp((z_off64_t)0, 0);
    uint32_t id17 = (uint32_t)test_x2nmodp((z_off64_t)0, 17);
    uint32_t id31 = (uint32_t)test_x2nmodp((z_off64_t)0, 31);
    uint32_t id63 = (uint32_t)test_x2nmodp((z_off64_t)0, 63);
    TEST_ASSERT_EQUAL_HEX32(0x80000000u, id0);
    TEST_ASSERT_EQUAL_HEX32(0x80000000u, id17);
    TEST_ASSERT_EQUAL_HEX32(0x80000000u, id31);
    TEST_ASSERT_EQUAL_HEX32(0x80000000u, id63);

    /* Small powers below degree 32 */
    uint32_t x1 = (uint32_t)test_x2nmodp((z_off64_t)1, 0);  /* x^1 */
    uint32_t x2 = (uint32_t)test_x2nmodp((z_off64_t)1, 1);  /* x^2 */
    uint32_t x3 = (uint32_t)test_x2nmodp((z_off64_t)1, 2);  /* x^4 */
    uint32_t x2_direct = (uint32_t)test_x2nmodp((z_off64_t)2, 0); /* x^2 */
    uint32_t x3_direct = (uint32_t)test_x2nmodp((z_off64_t)4, 0); /* x^4 */
    TEST_ASSERT_EQUAL_HEX32(0x40000000u, x1);
    TEST_ASSERT_EQUAL_HEX32(0x20000000u, x2);
    TEST_ASSERT_EQUAL_HEX32(0x10000000u, x3);
    TEST_ASSERT_EQUAL_HEX32(0x20000000u, x2_direct);
    TEST_ASSERT_EQUAL_HEX32(0x10000000u, x3_direct);
}

void test_x2nmodp_leftshift_equivalence(void) {
    /* x2nmodp(n, k) == x2nmodp(n << k, 0), choose values that won't overflow */
    z_off64_t n = (z_off64_t)0x1234;
    unsigned k = 5;
    z_off64_t n_shift = n << k; /* safe: 0x246800 */
    uint32_t a = (uint32_t)test_x2nmodp(n, k);
    uint32_t b = (uint32_t)test_x2nmodp(n_shift, 0);
    TEST_ASSERT_EQUAL_HEX32(b, a);

    n = 7;
    k = 3;
    n_shift = n << k; /* 56 */
    a = (uint32_t)test_x2nmodp(n, k);
    b = (uint32_t)test_x2nmodp(n_shift, 0);
    TEST_ASSERT_EQUAL_HEX32(b, a);
}

void test_x2nmodp_k_wraparound_32(void) {
    /* Periodicity: x^(n*2^(k+32)) == x^(n*2^k) */
    z_off64_t n = 7;
    uint32_t a = (uint32_t)test_x2nmodp(n, 3);
    uint32_t b = (uint32_t)test_x2nmodp(n, 3 + 32);
    uint32_t c = (uint32_t)test_x2nmodp(n, 3 + 64);
    TEST_ASSERT_EQUAL_HEX32(a, b);
    TEST_ASSERT_EQUAL_HEX32(a, c);

    n = 12345;
    a = (uint32_t)test_x2nmodp(n, 0);
    b = (uint32_t)test_x2nmodp(n, 32);
    TEST_ASSERT_EQUAL_HEX32(a, b);
}

void test_x2nmodp_matches_combine_gen64(void) {
    /* x2nmodp(n,3) must equal crc32_combine_gen64(n) */
    z_off64_t cases[] = {
        0, 1, 2, 3, 7, 31, 32, 33, 123, 1000000,
        (z_off64_t)1234567890123LL
    };
    size_t i;
    for (i = 0; i < sizeof(cases)/sizeof(cases[0]); i++) {
        z_off64_t n = cases[i];
        uint32_t a = (uint32_t)test_x2nmodp(n, 3);
        uint32_t b = (uint32_t)crc32_combine_gen64(n);
        TEST_ASSERT_EQUAL_HEX32(b, a);
    }

    /* Also test a very large n close to 2^63-1 */
    z_off64_t large_n = (z_off64_t)0x7fffffffffffffffLL;
    uint32_t a = (uint32_t)test_x2nmodp(large_n, 3);
    uint32_t b = (uint32_t)crc32_combine_gen64(large_n);
    TEST_ASSERT_EQUAL_HEX32(b, a);
}

void test_x2nmodp_additivity_via_multiplication(void) {
    /* x^( (a+b) * 2^k ) == x^(a*2^k) * x^(b*2^k) */
    unsigned k = 7;
    z_off64_t a_len = 5;
    z_off64_t b_len = 12;

    uint32_t op_a = (uint32_t)test_x2nmodp(a_len, k);
    uint32_t op_b = (uint32_t)test_x2nmodp(b_len, k);
    uint32_t op_ab = (uint32_t)test_x2nmodp(a_len + b_len, k);

    uint32_t prod = op_mul(op_a, op_b); /* multiply operators */
    TEST_ASSERT_EQUAL_HEX32(op_ab, prod);

    /* Associativity across three terms */
    z_off64_t c_len = 9;
    uint32_t op_c = (uint32_t)test_x2nmodp(c_len, k);
    uint32_t op_sum = (uint32_t)test_x2nmodp(a_len + b_len + c_len, k);
    uint32_t prod2 = op_mul(prod, op_c);
    TEST_ASSERT_EQUAL_HEX32(op_sum, prod2);
}

void test_x2nmodp_construct_from_bits(void) {
    /* Build x2nmodp(n, k) from its set bits: product over i in bits(n) of x2nmodp(1, k+i) */
    z_off64_t n = 13; /* 1101b */
    unsigned k = 7;

    uint32_t expected = (uint32_t)test_x2nmodp(n, k);

    uint32_t accum = (uint32_t)test_x2nmodp(0, k); /* identity x^0 */
    unsigned i = 0;
    z_off64_t tmp = n;
    while (tmp) {
        if (tmp & 1) {
            uint32_t term = (uint32_t)test_x2nmodp(1, k + i);
            accum = op_mul(accum, term);
        }
        tmp >>= 1;
        i++;
    }
    TEST_ASSERT_EQUAL_HEX32(expected, accum);
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_x2nmodp_identity_and_small_powers);
    RUN_TEST(test_x2nmodp_leftshift_equivalence);
    RUN_TEST(test_x2nmodp_k_wraparound_32);
    RUN_TEST(test_x2nmodp_matches_combine_gen64);
    RUN_TEST(test_x2nmodp_additivity_via_multiplication);
    RUN_TEST(test_x2nmodp_construct_from_bits);
    return UNITY_END();
}