import random
from math import ceil, log2

import pyrtl
from pyrtl import shift_left_logical, shift_right_logical
from pyrtl.rtllib import adders


def prng_lfsr(
    bitwidth: int,
    load: pyrtl.WireVector,
    req: pyrtl.WireVector,
    seed: pyrtl.WireVector = None,
) -> pyrtl.Register:
    """Builds a single-cycle PRNG using a 127 bits Fibonacci LFSR.

    A very fast and compact PRNG that generates a random number using only one clock
    cycle. Has a period of ``2**127 - 1``. Its linearity makes it a bit statistically
    weak, but it should be good enough for any noncryptographic purpose like test
    pattern generation.

    Example::

        load, req = pyrtl.Input(1, 'load'), pyrtl.Input(1, 'req')
        rand = pyrtl.Output(64, 'rand')

        rand <<= prngs.prng_lfsr(64, load, req)

        sim = pyrtl.Simulation()
        sim.step({'load': 1, 'req': 0}) # seed once at the beginning
        sim.step({'load': 0, 'req': 1})
        sim.step({'load': 0, 'req': 0})
        print(sim.inspect(rand))
        sim.tracer.render_trace(symbol_len=40, segment_size=1)

    Example with explicit seeding::

        seed =  pyrtl.Input(127, 'seed')
        load, req = pyrtl.Input(1, 'load'), pyrtl.Input(1, 'req')
        rand = pyrtl.Output(32, 'rand')

        rand <<= prngs.prng_lfsr(32, load, req, seed)

        sim = pyrtl.Simulation()
        sim.step({'load': 1, 'req': 0, 'seed': 0x102030405060708090a0b0c0d0e0f010})
        sim.step({'load': 0, 'req': 1, 'seed': 0x102030405060708090a0b0c0d0e0f010})
        sim.step({'load': 0, 'req': 0, 'seed': 0x102030405060708090a0b0c0d0e0f010})
        print(sim.inspect(rand))
        sim.tracer.render_trace(symbol_len=40, segment_size=1)

    :param bitwidth: The desired bitwidth of the random number.
    :param load: One bit signal to load the seed into the PRNG.
    :param req: One bit signal to request a random number.
    :param seed: 127 bits :class:`.WireVector`, defaults to ``None`` (self-seeding).
        Refrain from self-seeding if reseeding at run time is required.

    :return: :class:`.Register` containing the random number with the given
             ``bitwidth``.
    """
    # 127 bits is chosen because 127 is a mersenne prime, which makes the period of the
    # LFSR maximized at 2**127 - 1 for any requested bitwidth
    if seed is None:
        cryptogen = random.SystemRandom()
        seed = cryptogen.randrange(1, 2**127)  # seed itself if no seed signal is given

    lfsr = pyrtl.Register(max(bitwidth, 127))
    # leap ahead by shifting the LFSR bitwidth times
    leap_ahead = lfsr
    for _i in range(bitwidth):
        leap_ahead = pyrtl.concat(leap_ahead, leap_ahead[125] ^ leap_ahead[126])

    with pyrtl.conditional_assignment:
        with load:
            lfsr.next |= seed
        with req:
            lfsr.next |= leap_ahead
    return lfsr[:bitwidth]


def prng_xoroshiro128(
    bitwidth: int,
    load: pyrtl.WireVector,
    req: pyrtl.WireVector,
    seed: pyrtl.WireVector = None,
) -> tuple[pyrtl.WireVector, pyrtl.Register]:
    """Builds a PRNG using the Xoroshiro128+ algorithm in hardware.

    An efficient noncryptographic PRNG, has much smaller area than
    :func:`csprng_trivium`. But it does require a 64-bit adder to compute the output, so
    it is a bit slower. Has a period of ``2 ** 128 - 1``. Passes most statistical tests.
    Outputs a 64-bit random word each cycle, takes multiple cycles if more than 64 bits
    are requested, and MSBs of the random words are returned if the bitwidth is not a
    multiple of 64.

    See also http://xoroshiro.di.unimi.it/

    Example::

        load, req = pyrtl.Input(1, 'load'), pyrtl.Input(1, 'req')
        ready, rand = pyrtl.Output(1, 'ready'), pyrtl.Output(128, 'rand')

        ready_out, rand_out = prngs.prng_xoroshiro128(128, load, req)
        ready <<= ready_out
        rand <<= rand_out

        sim = pyrtl.Simulation()
        sim.step({'load': 1, 'req': 0})  # seed once at the beginning
        sim.step({'load': 0, 'req': 1})
        while sim.value[ready] == 0:  # or loop 2 cycles
            sim.step({'load': 0, 'req': 0})

        print(sim.inspect(rand))
        sim.tracer.render_trace(symbol_len=40, segment_size=1)

    :param bitwidth: The desired bitwidth of the random number.
    :param load: One bit signal to load the seed into the PRNG.
    :param req: One bit signal to request a random number.
    :param seed: 128 bits :class:`.WireVector`, defaults to ``None`` (self-seeding).
        Refrain from self-seeding if reseeding at run time is required.

    :return: ``(ready, rand)``, where ``ready`` is a one bit signal indicating that the
             random number has been produced, and ``rand`` is a :class:`.Register`
             containing the random number with the given ``bitwidth``.
    """
    if seed is None:
        cryptogen = random.SystemRandom()
        seed = cryptogen.randrange(1, 2**128)  # seed itself if no seed signal is given
    seed = pyrtl.as_wires(seed, 128)

    s0, s1 = (pyrtl.Register(64) for i in range(2))
    output = pyrtl.WireVector(64)
    # update internal states by xoring, rotating, and shifting
    _s1 = s0 ^ s1
    s0_next = (
        (shift_left_logical(s0, 55) | shift_right_logical(s0, 9))
        ^ shift_left_logical(_s1, 14)
        ^ _s1
    )
    s1_next = shift_left_logical(_s1, 36) | shift_right_logical(_s1, 28)
    output <<= adders.kogge_stone(s0, s1)

    gen_cycles = ceil(bitwidth / 64)
    counter_bitwidth = ceil(log2(gen_cycles)) if gen_cycles > 1 else 1
    rand = pyrtl.Register(gen_cycles * 64)
    counter = pyrtl.Register(counter_bitwidth, "counter")
    gen_done = counter == gen_cycles - 1
    state = pyrtl.Register(1)
    WAIT, GEN = (pyrtl.Const(x) for x in range(2))
    with pyrtl.conditional_assignment:
        with load:
            s0.next |= seed[:64]
            s1.next |= seed[64:]
            state.next |= WAIT
        with req:
            counter.next |= 0
            s0.next |= s0_next
            s1.next |= s1_next
            rand.next |= pyrtl.concat(rand, output)
            state.next |= GEN
        with state == GEN:
            with ~gen_done:
                counter.next |= counter + 1
                s0.next |= s0_next
                s1.next |= s1_next
                rand.next |= pyrtl.concat(rand, output)

    ready = ~load & ~req & (state == GEN) & gen_done
    return ready, rand[-bitwidth:]  # return MSBs because LSBs are less random


def csprng_trivium(
    bitwidth: int,
    load: pyrtl.WireVector,
    req: pyrtl.WireVector,
    seed: pyrtl.WireVector = None,
    bits_per_cycle: int = 64,
) -> tuple[pyrtl.WireVector, pyrtl.Register]:
    """Builds a cyptographically secure PRNG using the Trivium stream cipher.

    This PRNG uses Trivium's key stream as its random bits output. Both ``seed`` and the
    key stream are MSB first (the earliest bit is stored at the MSB). Trivium has a seed
    initialization stage that discards the first weak 1152 output bits after each
    loading. Generation stage can take multiple cycles as well depending on the given
    bitwidth and ``bits_per_cycle``. Has smaller gate area and faster speed than AES-CTR
    and any other stream cipher. Passes all known statistical tests. Can be used to
    generate encryption keys or IVs. Designed to securely generate up to ``2 ** 64``
    bits. If more than ``2 ** 64`` bits are needed, must reseed after each generation of
    ``2 ** 64`` bits.

    Trivium specifications: http://www.ecrypt.eu.org/stream/ciphers/trivium/trivium.pdf

    See also the eSTREAM portfolio page: http://www.ecrypt.eu.org/stream/e2-trivium.html

    Example::

        load, req = pyrtl.Input(1, 'load'), pyrtl.Input(1, 'req')
        ready, rand = pyrtl.Output(1, 'ready'), pyrtl.Output(128, 'rand')

        ready_out, rand_out = prngs.csprng_trivium(128, load, req)
        ready <<= ready_out
        rand <<= rand_out

        sim = pyrtl.Simulation()
        # Seed only in the first cycle.
        sim.step({'load': 1, 'req': 0})
        while sim.value[ready] == 0:  # or loop 19 cycles
            sim.step({'load': 0, 'req': 0})

        sim.step({'load': 0, 'req': 1})
        while sim.value[ready] == 0:  # or loop 2 cycles
            sim.step({'load': 0, 'req': 0})

        print(sim.inspect(rand))
        sim.tracer.render_trace(symbol_len=45, segment_size=5)

    :param bitwidth: The desired bitwidth of the random number.
    :param load: One bit signal to load the seed into the PRNG
    :param req: One bit signal to request a random number.
    :param seed: 160 bits :class:`.WireVector` (80 bits key + 80 bits IV), defaults to
        ``None`` (self-seeding). Refrain from self-seeding if reseeding at run time is
        needed.
    :param bits_per_cycle: The number of output bits to generate in parallel each cycle,
        up to ``64`` bits. Must be a power of two, so ``bits_per_cycle`` must be one of
        ``1``, ``2``, ``4``, ``8``, ``16``, ``32``, or ``64``.

    :return: ``(ready, rand)``, where ``ready`` is a one bit signal indicating that the
             random number has been produced or the seed has been initialized, and
             ``rand`` is a :class:`.Register` containing the random number with the
             given ``bitwidth``.
    """
    if (64 // bits_per_cycle) * bits_per_cycle != 64:
        msg = "bits_per_cycle is invalid"
        raise pyrtl.PyrtlError(msg)
    if seed is None:
        cryptogen = random.SystemRandom()
        seed = cryptogen.randrange(2**160)  # seed itself if no seed signal is given
    seed = pyrtl.as_wires(seed, 160)
    key = seed[80:]
    iv = seed[:80]

    a = pyrtl.Register(93)
    b = pyrtl.Register(84)
    c = pyrtl.Register(111)
    feedback_a, feedback_b, feedback_c, output = ([] for i in range(4))
    for i in range(bits_per_cycle):
        t1 = a[65 - i] ^ a[92 - i]
        t2 = b[68 - i] ^ b[83 - i]
        t3 = c[65 - i] ^ c[110 - i]
        feedback_a.append(t3 ^ c[108 - i] & c[109 - i] ^ a[68 - i])
        feedback_b.append(t1 ^ a[90 - i] & a[91 - i] ^ b[77 - i])
        feedback_c.append(t2 ^ b[81 - i] & b[82 - i] ^ c[86 - i])
        output.append(t1 ^ t2 ^ t3)
    # update internal states by shifting bits_per_cycle times
    a_next = pyrtl.concat(a, *feedback_a)
    b_next = pyrtl.concat(b, *feedback_b)
    c_next = pyrtl.concat(c, *feedback_c)

    init_cycles = 1152 // bits_per_cycle
    gen_cycles = ceil(bitwidth / bits_per_cycle)
    counter_bitwidth = ceil(log2(max(init_cycles + 1, gen_cycles)))
    rand = pyrtl.Register(bitwidth)
    counter = pyrtl.Register(counter_bitwidth, "counter")
    init_done = counter == init_cycles
    gen_done = counter == gen_cycles - 1
    state = pyrtl.Register(2)
    WAIT, INIT, GEN = (pyrtl.Const(x) for x in range(3))
    with pyrtl.conditional_assignment:
        with load:
            counter.next |= 0
            a.next |= key
            b.next |= iv
            c.next |= pyrtl.concat(pyrtl.Const("3'b111"), pyrtl.Const(0, 108))
            state.next |= INIT
        with req:
            counter.next |= 0
            a.next |= a_next
            b.next |= b_next
            c.next |= c_next
            rand.next |= pyrtl.concat(rand, *output)
            state.next |= GEN
        with state == INIT:
            with ~init_done:
                counter.next |= counter + 1
                a.next |= a_next
                b.next |= b_next
                c.next |= c_next
        with state == GEN:
            with ~gen_done:
                counter.next |= counter + 1
                a.next |= a_next
                b.next |= b_next
                c.next |= c_next
                rand.next |= pyrtl.concat(rand, *output)

    ready = ~load & ~req & ((state == INIT) & init_done | (state == GEN) & gen_done)
    return ready, rand
