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

/**
 * @brief Calculates (a * b) % c taking into account that a * b might overflow
 *
 * @param a First factor
 * @param b Second factor
 * @param mod The modulo
 * @return Resulting value
 */
uint64_t mulmod(uint64_t a, uint64_t b, uint64_t mod);

/**
 * @brief Modular exponentiation
 *
 * @param base
 * @param exponent
 * @param mod
 * @return
 */
uint64_t modulo(uint64_t base, uint64_t exponent, uint64_t mod);

/**
 * @brief Miller-Rabin Primality Test
 *
 * This function performs the Miller-Rabin primality test to determine whether
 * a given number is prime. The number of iterations increases the accuracy of
 * the test but also increases computation time.
 *
 * @param p The number to test for primality (must be a positive integer).
 * @param iteration The number of test rounds to perform (higher means more
 * accuracy).
 * @return 1 if the number is probably prime, 0 if it is composite.
 */
uint32_t miller_rabin(uint64_t p, uint32_t iteration);

/* Below code is source, not header */
/************************************/

uint64_t mulmod(uint64_t a, uint64_t b, uint64_t mod) {
    uint64_t x = 0, y = a % mod;

    while (b > 0) {
        if (b % 2 == 1) {
            x = (x + y) % mod;
        }
        y = (y * 2) % mod;
        b /= 2;
    }

    return x % mod;
}

uint64_t modulo(uint64_t base, uint64_t exponent, uint64_t mod) {
    uint64_t x = 1;
    uint64_t y = base;

    while (exponent > 0) {
        if (exponent % 2 == 1)
            x = (x * y) % mod;

        y = (y * y) % mod;

        exponent = exponent / 2;
    }

    return x % mod;
}

uint32_t miller_rabin(uint64_t p, uint32_t iteration) {
    uint32_t i;
    uint64_t s;

    if (p < 2) {
        return 0;
    }

    if (p != 2 && p % 2 == 0) {
        return 0;
    }

    s = p - 1;

    while (s % 2 == 0) {
        s /= 2;
    }

    for (i = 0; i < iteration; i++) {
        uint64_t a = rand() % (p - 1) + 1, temp = s;
        uint64_t mod = modulo(a, temp, p);

        while (temp != p - 1 && mod != 1 && mod != p - 1) {
            mod = mulmod(mod, mod, p);
            temp *= 2;
        }

        if (mod != p - 1 && temp % 2 == 0) {
            return 0;
        }
    }
    return 1;
}