#include "../../unity/unity.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

/* Unity requires these even if empty */
void setUp(void) {
  /* Setup code here, or leave empty */
}
void tearDown(void) {
  /* Cleanup code here, or leave empty */
}

/* Helper: find the index of a given small prime in primes_ptab.
   Returns -1 if not found (should not happen for small primes we use). */
static int find_prime_index(mp_limb_t prime)
{
  for (int i = 0; i < PRIMES_PTAB_ENTRIES; i++)
    {
      if ((mp_limb_t)primes_ptab[i] == prime)
        return i;
    }
  return -1;
}

/* Helper: compute the base index and offset (ioff) for a given index,
   matching the 8-way unrolling the kernel supports. */
static void base_and_offset_for_index(int idx, int *base, int *ioff)
{
  int off = idx % 8;
  *ioff = off;
  *base = idx - off;
}

static void init_factors(struct factors *f)
{
  memset(f, 0, sizeof(*f));
}

void test_divblock_divides_once_by_3(void)
{
  int idx = find_prime_index(3);
  TEST_ASSERT_NOT_EQUAL(-1, idx);
  int base = 0, off = 0;
  base_and_offset_for_index(idx, &base, &off);
  TEST_ASSERT_EQUAL_INT(idx, base + off);

  struct factors f; init_factors(&f);
  mp_limb_t t0 = (mp_limb_t)3;
  mp_limb_t r = divblock(&f, t0, &primes_dtab[base], (idx_t)base, off);

  TEST_ASSERT(r == (mp_limb_t)1);
  TEST_ASSERT(f.nfactors == 1);
  TEST_ASSERT(f.p[0] == (mp_limb_t)3);
  TEST_ASSERT(f.e[0] == 1);
}

void test_divblock_divides_twice_by_3(void)
{
  int idx = find_prime_index(3);
  TEST_ASSERT_NOT_EQUAL(-1, idx);
  int base = 0, off = 0;
  base_and_offset_for_index(idx, &base, &off);

  struct factors f; init_factors(&f);
  mp_limb_t t0 = (mp_limb_t)9; /* 3^2 */
  mp_limb_t r = divblock(&f, t0, &primes_dtab[base], (idx_t)base, off);

  TEST_ASSERT(r == (mp_limb_t)1);
  TEST_ASSERT(f.nfactors == 1);
  TEST_ASSERT(f.p[0] == (mp_limb_t)3);
  TEST_ASSERT(f.e[0] == 2);
}

void test_divblock_non_divisible_returns_input(void)
{
  int idx = find_prime_index(3);
  TEST_ASSERT_NOT_EQUAL(-1, idx);
  int base = 0, off = 0;
  base_and_offset_for_index(idx, &base, &off);

  struct factors f; init_factors(&f);
  mp_limb_t t0 = (mp_limb_t)10; /* Not divisible by 3 */
  mp_limb_t r = divblock(&f, t0, &primes_dtab[base], (idx_t)base, off);

  TEST_ASSERT(r == t0);
  TEST_ASSERT(f.nfactors == 0);
}

void test_divblock_returns_remaining_quotient_after_divisions(void)
{
  int idx = find_prime_index(3);
  TEST_ASSERT_NOT_EQUAL(-1, idx);
  int base = 0, off = 0;
  base_and_offset_for_index(idx, &base, &off);

  struct factors f; init_factors(&f);
  /* t0 = 3^3 * 10 = 270, should divide out 3 three times and return 10 */
  mp_limb_t t0 = (mp_limb_t)27 * (mp_limb_t)10;
  mp_limb_t r = divblock(&f, t0, &primes_dtab[base], (idx_t)base, off);

  TEST_ASSERT(r == (mp_limb_t)10);
  TEST_ASSERT(f.nfactors == 1);
  TEST_ASSERT(f.p[0] == (mp_limb_t)3);
  TEST_ASSERT(f.e[0] == 3);
}

void test_divblock_works_with_nonzero_offset_prime_23(void)
{
  int idx = find_prime_index(23);
  TEST_ASSERT_NOT_EQUAL(-1, idx);
  int base = 0, off = 0;
  base_and_offset_for_index(idx, &base, &off);
  TEST_ASSERT(off >= 0 && off <= 7);
  TEST_ASSERT_EQUAL_INT(idx, base + off);

  struct factors f; init_factors(&f);
  mp_limb_t t0 = (mp_limb_t)23 * (mp_limb_t)7; /* 161 */
  mp_limb_t r = divblock(&f, t0, &primes_dtab[base], (idx_t)base, off);

  TEST_ASSERT(r == (mp_limb_t)7);
  TEST_ASSERT(f.nfactors == 1);
  TEST_ASSERT(f.p[0] == (mp_limb_t)23);
  TEST_ASSERT(f.e[0] == 1);
}

void test_divblock_increments_existing_factor_multiplicity(void)
{
  int idx = find_prime_index(5);
  TEST_ASSERT_NOT_EQUAL(-1, idx);
  int base = 0, off = 0;
  base_and_offset_for_index(idx, &base, &off);

  struct factors f; init_factors(&f);
  /* Pre-populate factors with prime 5 multiplicity 2 */
  f.nfactors = 1;
  f.p[0] = (mp_limb_t)5;
  f.e[0] = 2;

  mp_limb_t t0 = (mp_limb_t)5;
  mp_limb_t r = divblock(&f, t0, &primes_dtab[base], (idx_t)base, off);

  TEST_ASSERT(r == (mp_limb_t)1);
  TEST_ASSERT(f.nfactors == 1);
  TEST_ASSERT(f.p[0] == (mp_limb_t)5);
  TEST_ASSERT(f.e[0] == 3);
}

void test_divblock_t0_equal_1_no_change(void)
{
  int idx = find_prime_index(7);
  TEST_ASSERT_NOT_EQUAL(-1, idx);
  int base = 0, off = 0;
  base_and_offset_for_index(idx, &base, &off);

  struct factors f; init_factors(&f);
  mp_limb_t t0 = (mp_limb_t)1;
  mp_limb_t r = divblock(&f, t0, &primes_dtab[base], (idx_t)base, off);

  TEST_ASSERT(r == (mp_limb_t)1);
  TEST_ASSERT(f.nfactors == 0);
}

int main(void)
{
  UNITY_BEGIN();
  RUN_TEST(test_divblock_divides_once_by_3);
  RUN_TEST(test_divblock_divides_twice_by_3);
  RUN_TEST(test_divblock_non_divisible_returns_input);
  RUN_TEST(test_divblock_returns_remaining_quotient_after_divisions);
  RUN_TEST(test_divblock_works_with_nonzero_offset_prime_23);
  RUN_TEST(test_divblock_increments_existing_factor_multiplicity);
  RUN_TEST(test_divblock_t0_equal_1_no_change);
  return UNITY_END();
}