#include "../../unity/unity.h"
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>
#include <gmp.h>

/* We rely on symbols from the program under test being in the same
   translation unit due to direct inclusion. In particular, we use:
   - W_TYPE_SIZE, mp_limb_t
   - uuint, make_uuint, hi(), lo()
   - the target function: mod2(a1, a0, d1, d0)
*/

/* Helper: reference implementation using GMP for (a1,a0) mod (d1,d0). */
static uuint ref_mod2_gmp(mp_limb_t a1, mp_limb_t a0, mp_limb_t d1, mp_limb_t d0)
{
    /* Build A and D as mpz from two limbs (least significant limb first). */
    mpz_t A, D, R;
    mpz_init(A);
    mpz_init(D);
    mpz_init(R);

    mp_limb_t a_arr[2];
    mp_limb_t d_arr[2];
    a_arr[0] = a0; a_arr[1] = a1;
    d_arr[0] = d0; d_arr[1] = d1;

    mpz_import(A, 2, -1, sizeof(mp_limb_t), 0, 0, a_arr);
    mpz_import(D, 2, -1, sizeof(mp_limb_t), 0, 0, d_arr);

    mpz_mod(R, A, D);

    mp_limb_t r_arr[2] = {0, 0};
    size_t count = 0;
    mpz_export(r_arr, &count, -1, sizeof(mp_limb_t), 0, 0, R);

    mp_limb_t r0 = (count >= 1) ? r_arr[0] : (mp_limb_t)0;
    mp_limb_t r1 = (count >= 2) ? r_arr[1] : (mp_limb_t)0;

    mpz_clear(A);
    mpz_clear(D);
    mpz_clear(R);

    return make_uuint(r1, r0);
}

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

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

/* Test: modulus by B (d = (1,0)) should return low word. */
void test_mod2_mod_by_base_returns_low_word(void)
{
    mp_limb_t a1 = (((mp_limb_t)1) << (W_TYPE_SIZE - 1)) | (mp_limb_t)123u; /* ensure nonzero high */
    mp_limb_t a0 = (mp_limb_t)0xDEADBEEFu;
    mp_limb_t d1 = (mp_limb_t)1;
    mp_limb_t d0 = (mp_limb_t)0;

    uuint r = mod2(a1, a0, d1, d0);

    TEST_ASSERT_TRUE(hi(r) == 0);
    TEST_ASSERT_TRUE(lo(r) == a0);
}

/* Test: if A < D, result should be A (case with a1 nonzero). */
void test_mod2_A_less_than_D_returns_A(void)
{
    mp_limb_t a1 = (mp_limb_t)1;
    mp_limb_t a0 = (mp_limb_t)12345u;
    mp_limb_t d1 = (mp_limb_t)2; /* D > A */
    mp_limb_t d0 = (mp_limb_t)0;

    uuint r = mod2(a1, a0, d1, d0);

    TEST_ASSERT_TRUE(hi(r) == a1);
    TEST_ASSERT_TRUE(lo(r) == a0);
}

/* Test: if a1 == 0 and d1 != 0, A < D, so remainder is A. */
void test_mod2_a1_zero_returns_a(void)
{
    mp_limb_t a1 = (mp_limb_t)0;
    mp_limb_t a0 = (mp_limb_t)987u;
    mp_limb_t d1 = (mp_limb_t)1; /* any nonzero */
    mp_limb_t d0 = (mp_limb_t)123u;

    uuint r = mod2(a1, a0, d1, d0);

    TEST_ASSERT_TRUE(hi(r) == 0);
    TEST_ASSERT_TRUE(lo(r) == a0);
}

/* Test: general cases compared against GMP reference, ensuring cnt > 0 path exercised. */
void test_mod2_general_against_gmp(void)
{
    /* Case 1: Large a1, small d1 to ensure cnt > 0, nonzero d0. */
    {
        mp_limb_t a1 = (((mp_limb_t)1) << (W_TYPE_SIZE - 1)) | (mp_limb_t)0x55u;
        mp_limb_t a0 = (mp_limb_t)0xA5A5A5A5u;
        mp_limb_t d1 = a1 >> 5; /* strictly smaller top bit, ensures cnt > 0 */
        if (d1 == 0) d1 = 1;    /* ensure d1 != 0 */
        mp_limb_t d0 = (mp_limb_t)0x54321u;

        uuint r = mod2(a1, a0, d1, d0);
        uuint g = ref_mod2_gmp(a1, a0, d1, d0);
        TEST_ASSERT_TRUE(hi(r) == hi(g));
        TEST_ASSERT_TRUE(lo(r) == lo(g));
    }

    /* Case 2: Another configuration with different shift. */
    {
        mp_limb_t a1 = ((((mp_limb_t)1) << (W_TYPE_SIZE - 2)) | (mp_limb_t)77u);
        mp_limb_t a0 = (mp_limb_t)0x12345678u;
        mp_limb_t d1 = a1 >> 3; /* ensure cnt > 0 */
        if (d1 == 0) d1 = 1;
        mp_limb_t d0 = (mp_limb_t)0x9ABCDEFFu;

        uuint r = mod2(a1, a0, d1, d0);
        uuint g = ref_mod2_gmp(a1, a0, d1, d0);
        TEST_ASSERT_TRUE(hi(r) == hi(g));
        TEST_ASSERT_TRUE(lo(r) == lo(g));
    }

    /* Case 3: d0 = 0, still with cnt > 0. */
    {
        mp_limb_t a1 = (((mp_limb_t)1) << (W_TYPE_SIZE - 1));
        mp_limb_t a0 = (mp_limb_t)0xCAFEBABEu;
        mp_limb_t d1 = a1 >> 7; /* ensure cnt > 0 */
        if (d1 == 0) d1 = 1;
        mp_limb_t d0 = (mp_limb_t)0;

        uuint r = mod2(a1, a0, d1, d0);
        uuint g = ref_mod2_gmp(a1, a0, d1, d0);
        TEST_ASSERT_TRUE(hi(r) == hi(g));
        TEST_ASSERT_TRUE(lo(r) == lo(g));
    }

    /* Case 4: very small d1 to make many iterations, nonzero d0. */
    {
        mp_limb_t a1 = (((mp_limb_t)1) << (W_TYPE_SIZE - 1)) | (mp_limb_t)0xF0u;
        mp_limb_t a0 = (mp_limb_t)0x0BADF00Du;
        mp_limb_t d1 = (mp_limb_t)1;
        mp_limb_t d0 = (mp_limb_t)0x11111111u;

        uuint r = mod2(a1, a0, d1, d0);
        uuint g = ref_mod2_gmp(a1, a0, d1, d0);
        TEST_ASSERT_TRUE(hi(r) == hi(g));
        TEST_ASSERT_TRUE(lo(r) == lo(g));
    }

    /* Case 5: large values, random-looking, ensure cnt > 0 by shrinking d1. */
    {
        mp_limb_t a1 = (mp_limb_t)(~(mp_limb_t)0); /* MP_LIMB_MAX */
        mp_limb_t a0 = (mp_limb_t)0xFFFFFFFFu;
        mp_limb_t d1 = a1 >> 4; /* ensure cnt > 0 */
        if (d1 == 0) d1 = 1;
        mp_limb_t d0 = (mp_limb_t)0xFEEDFACEu;

        uuint r = mod2(a1, a0, d1, d0);
        uuint g = ref_mod2_gmp(a1, a0, d1, d0);
        TEST_ASSERT_TRUE(hi(r) == hi(g));
        TEST_ASSERT_TRUE(lo(r) == lo(g));
    }
}

/* Test: divisible case yielding zero remainder with D = (1, d0), A = 2*D, avoiding overflow. */
void test_mod2_divisible_case_zero_remainder(void)
{
    mp_limb_t d1 = (mp_limb_t)1;
    mp_limb_t d0 = (mp_limb_t)0x11111111u;

    /* Compute A = 2 * D without overflowing beyond two words. */
    mp_limb_t a0 = d0 << 1;
    mp_limb_t carry = (d0 >> (W_TYPE_SIZE - 1));
    mp_limb_t a1 = (d1 << 1) + carry; /* With d1=1, this can't overflow one limb. */

    uuint r = mod2(a1, a0, d1, d0);

    TEST_ASSERT_TRUE(hi(r) == 0);
    TEST_ASSERT_TRUE(lo(r) == 0);
}

int main(void)
{
    UNITY_BEGIN();
    RUN_TEST(test_mod2_mod_by_base_returns_low_word);
    RUN_TEST(test_mod2_A_less_than_D_returns_A);
    RUN_TEST(test_mod2_a1_zero_returns_a);
    RUN_TEST(test_mod2_general_against_gmp);
    RUN_TEST(test_mod2_divisible_case_zero_remainder);
    return UNITY_END();
}