from enum import Enum, auto
from random import getrandbits

from fixedint import FixedInt

# p and q gotta be primes that are congruent to 3 mod 4
p = 30000000091
q = 40000000003
m = p*q

class PRNGAlgorithm(Enum):
    XORSHIFT = auto()
    BBS = auto()

# Sometimes the number returned by getrandbits()
# doesn't have exactly the number of bits requested,
# this function exits to assure that it will.
def genSeed(nBits: int) -> int:
    s = getrandbits(nBits)

    while s.bit_length() != nBits:
        s = getrandbits(nBits)

    return FixedInt(nBits, signed=False, mutable=True)(s)

def xorshift(nBits: int, s: FixedInt) -> int:
    """Return a random number with nBits of length generated]
    by the Xorshift algorithm.

    Args:
        nBits: The number of bits desired
        s: The seed for the PRNG
    """

    while True:
        s ^= (s << 13)
        s ^= (s >> 7)
        s ^= (s << 17)

        # Making sure the number has the correct number of bits
        if int(s).bit_length() == nBits:
            break

    return s

def bbs(nBits: int, seed: int) -> int:
    """Return a random number with nBits of length generated by
    Blum-Blum-Shub algorithm.

    Generate numbers using the formula xn = x(n-1)² mod m and
    build a new number using the LSB of each generated number.

    Args:
        nBits: The number of bits desired
        s: The seed for the PRNG
    """

    b = 0

    x = seed ** 2 % m

    b = b | (x & 1)

    for i in range(1, nBits):
        x = x ** 2 % m

        b = b | ((x & 1) << i)

    return b

def randomNBits(nBits: int, seed: int, alg=PRNGAlgorithm.XORSHIFT) -> int:
    """Return a random number with nBits of length.

    Args:
        nBits: The number of bits desired
        seed: The seed for the PRNG
        alg: The algorithm desired for generating the number
    """

    if alg == PRNGAlgorithm.XORSHIFT:

        return xorshift(nBits, seed)
    elif alg == PRNGAlgorithm.BBS:
        seed = genSeed(nBits)

        return bbs(nBits, seed)
