#include "unity/unity.h"
#include "zlib.h"

#include <stdint.h>
#include <stdlib.h>

/* Match the module's selection of word width and type for the wrapper. */
#if defined(__x86_64__) || defined(__aarch64__)
  #define TEST_W 8
  typedef unsigned long long test_word_t;
  extern unsigned long long test_byte_swap(unsigned long long word);
#else
  #define TEST_W 4
  typedef unsigned long test_word_t;
  extern unsigned long test_byte_swap(unsigned long word);
#endif

/* Helper: reference byte-swap independent of the module under test. */
static test_word_t ref_swap(test_word_t v) {
    test_word_t out = 0;
    for (int i = 0; i < TEST_W; i++) {
        unsigned int b = (unsigned int)((v >> (8 * i)) & (test_word_t)0xFF);
        out |= ((test_word_t)b) << (8 * (TEST_W - 1 - i));
    }
    return out;
}

void setUp(void) {
    /* Setup code here, or leave empty */
}

void tearDown(void) {
    /* Cleanup code here, or leave empty */
}

/* Test: zero should remain zero after swapping. */
void test_byte_swap_zero(void) {
    test_word_t in = (test_word_t)0;
    test_word_t out = test_byte_swap(in);
    TEST_ASSERT_EQUAL_HEX64((uint64_t)0, (uint64_t)out);
}

/* Test: known pattern reversal. */
void test_byte_swap_known_pattern(void) {
#if TEST_W == 4
    test_word_t in = (test_word_t)0x01234567UL;
    test_word_t expected = (test_word_t)0x67452301UL;
    test_word_t out = test_byte_swap(in);
    TEST_ASSERT_EQUAL_HEX32((uint32_t)expected, (uint32_t)out);
#else
    test_word_t in = (test_word_t)0x0123456789ABCDEFULL;
    test_word_t expected = (test_word_t)0xEFCDAB8967452301ULL;
    test_word_t out = test_byte_swap(in);
    /* Use 64-bit comparator to avoid truncation */
    TEST_ASSERT_EQUAL_HEX64((uint64_t)expected, (uint64_t)out);
#endif
}

/* Test: all bytes same remain unchanged after reversal. */
void test_byte_swap_all_bytes_same(void) {
#if TEST_W == 4
    test_word_t in1 = (test_word_t)0x01010101UL;
    test_word_t in2 = (test_word_t)0xFFFFFFFFUL;
    TEST_ASSERT_EQUAL_HEX32((uint32_t)in1, (uint32_t)test_byte_swap(in1));
    TEST_ASSERT_EQUAL_HEX32((uint32_t)in2, (uint32_t)test_byte_swap(in2));
#else
    test_word_t in1 = (test_word_t)0x0101010101010101ULL;
    test_word_t in2 = (test_word_t)0xFFFFFFFFFFFFFFFFULL;
    TEST_ASSERT_EQUAL_HEX64((uint64_t)in1, (uint64_t)test_byte_swap(in1));
    TEST_ASSERT_EQUAL_HEX64((uint64_t)in2, (uint64_t)test_byte_swap(in2));
#endif
}

/* Test: applying byte_swap twice yields the original value (involution property). */
void test_byte_swap_involution(void) {
    /* A small but diverse set of inputs */
#if TEST_W == 4
    test_word_t samples[] = {
        0x00000000UL, 0x00000001UL, 0x80000000UL, 0x7FFFFFFFUL,
        0x01234567UL, 0x89ABCDEFUL, 0xA5A5A5A5UL, 0x5A5A5A5AUL
    };
    size_t count = sizeof(samples) / sizeof(samples[0]);
    for (size_t i = 0; i < count; i++) {
        test_word_t x = samples[i];
        test_word_t y = test_byte_swap(test_byte_swap(x));
        TEST_ASSERT_EQUAL_HEX32((uint32_t)x, (uint32_t)y);
    }
#else
    test_word_t samples[] = {
        0x0000000000000000ULL, 0x0000000000000001ULL,
        0x8000000000000000ULL, 0x7FFFFFFFFFFFFFFFULL,
        0x0123456789ABCDEFULL, 0xFEDCBA9876543210ULL,
        0xA5A5A5A5A5A5A5A5ULL, 0x5A5A5A5A5A5A5A5AULL
    };
    size_t count = sizeof(samples) / sizeof(samples[0]);
    for (size_t i = 0; i < count; i++) {
        test_word_t x = samples[i];
        test_word_t y = test_byte_swap(test_byte_swap(x));
        TEST_ASSERT_EQUAL_HEX64((uint64_t)x, (uint64_t)y);
    }
#endif
}

/* Test: a single non-zero byte moves to the mirrored position. */
void test_byte_swap_single_byte_moves(void) {
    for (int pos = 0; pos < TEST_W; pos++) {
        test_word_t in = 0;
        /* Place 0xAA at byte position pos (0 = least-significant byte) */
        in |= ((test_word_t)0xAA) << (8 * pos);
        test_word_t out = test_byte_swap(in);
        test_word_t expected = ((test_word_t)0xAA) << (8 * (TEST_W - 1 - pos));
#if TEST_W == 4
        TEST_ASSERT_EQUAL_HEX32((uint32_t)expected, (uint32_t)out);
#else
        TEST_ASSERT_EQUAL_HEX64((uint64_t)expected, (uint64_t)out);
#endif
    }
}

/* Test: palindromic-bytes remain unchanged after swap. */
void test_byte_swap_palindromic_bytes(void) {
#if TEST_W == 4
    /* bytes: [0x12, 0x34, 0x34, 0x12] */
    test_word_t in = ((test_word_t)0x12 << 24) |
                     ((test_word_t)0x34 << 16) |
                     ((test_word_t)0x34 << 8)  |
                     ((test_word_t)0x12);
    test_word_t out = test_byte_swap(in);
    TEST_ASSERT_EQUAL_HEX32((uint32_t)in, (uint32_t)out);
#else
    /* bytes: [0x01,0x23,0x45,0x67,0x67,0x45,0x23,0x01] */
    test_word_t in = ((test_word_t)0x01ULL << 56) |
                     ((test_word_t)0x23ULL << 48) |
                     ((test_word_t)0x45ULL << 40) |
                     ((test_word_t)0x67ULL << 32) |
                     ((test_word_t)0x67ULL << 24) |
                     ((test_word_t)0x45ULL << 16) |
                     ((test_word_t)0x23ULL << 8)  |
                     ((test_word_t)0x01ULL);
    test_word_t out = test_byte_swap(in);
    TEST_ASSERT_EQUAL_HEX64((uint64_t)in, (uint64_t)out);
#endif
}

/* Test: compare against a reference swap for several generated patterns. */
void test_byte_swap_against_reference(void) {
    /* Generate patterns where byte i = (seed + i) & 0xFF for a few seeds. */
    for (unsigned seed = 0; seed < 8; seed++) {
        test_word_t in = 0;
        for (int i = 0; i < TEST_W; i++) {
            test_word_t b = (test_word_t)((seed + i * 37) & 0xFF);
            in |= b << (8 * i); /* little-endian composition for input */
        }
        test_word_t expected = ref_swap(in);
        test_word_t out = test_byte_swap(in);
#if TEST_W == 4
        TEST_ASSERT_EQUAL_HEX32((uint32_t)expected, (uint32_t)out);
#else
        TEST_ASSERT_EQUAL_HEX64((uint64_t)expected, (uint64_t)out);
#endif
    }
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_byte_swap_zero);
    RUN_TEST(test_byte_swap_known_pattern);
    RUN_TEST(test_byte_swap_all_bytes_same);
    RUN_TEST(test_byte_swap_involution);
    RUN_TEST(test_byte_swap_single_byte_moves);
    RUN_TEST(test_byte_swap_palindromic_bytes);
    RUN_TEST(test_byte_swap_against_reference);
    return UNITY_END();
}