#!/usr/bin/env python3
#
# Copyright (c) The mldsa-native project authors
# SPDX-License-Identifier: LicenseRef-PD-hp OR CC0-1.0 OR 0BSD OR MIT-0 OR MIT

# References
# ==========
#
# - [surf]
#   SURF: Simple Unpredictable Random Function
#   Daniel J. Bernstein
#   https://cr.yp.to/papers.html#surf

# Based on @[surf].

# WARNING
#
# The randombytes() implementation in this file is for TESTING ONLY.
# You MUST NOT use this implementation outside of testing.

import sys

MASK32 = 0xFFFFFFFF

seed = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9, 3,
        2, 3, 8, 4, 6, 2, 6, 4, 3, 3, 8, 3, 2, 7, 9, 5]

state_in = [0] * 12
out = [0] * 8
outleft = 0


def randombytes_reset():
    global state_in, out, outleft
    state_in = [0] * 12
    out = [0] * 8
    outleft = 0


def _rotate(x, b):
    return ((x << b) | (x >> (32 - b))) & MASK32


def _surf():
    global out
    t = [(state_in[i] ^ seed[12 + i]) & MASK32 for i in range(12)]
    out = [seed[24 + i] for i in range(8)]
    x = t[11]
    sum_val = 0
    for _ in range(2):
        for _ in range(16):
            sum_val = (sum_val + 0x9E3779B9) & MASK32
            for i, b in [(0, 5), (1, 7), (2, 9), (3, 13),
                         (4, 5), (5, 7), (6, 9), (7, 13),
                         (8, 5), (9, 7), (10, 9), (11, 13)]:
                val = (((x ^ seed[i]) + sum_val) & MASK32) ^ _rotate(x, b)
                t[i] = (t[i] + val) & MASK32
                x = t[i]
        for i in range(8):
            out[i] = (out[i] ^ t[i + 4]) & MASK32


def randombytes(n):
    global outleft
    buf = []
    for _ in range(n):
        if not outleft:
            state_in[0] = (state_in[0] + 1) & MASK32
            if state_in[0] == 0:
                state_in[1] = (state_in[1] + 1) & MASK32
                if state_in[1] == 0:
                    state_in[2] = (state_in[2] + 1) & MASK32
                    if state_in[2] == 0:
                        state_in[3] = (state_in[3] + 1) & MASK32
            _surf()
            outleft = 8
        outleft -= 1
        buf.append(out[outleft] & 0xFF)
    return bytes(buf)


if __name__ == "__main__":
    num_bytes = int(sys.argv[1]) if len(sys.argv) > 1 else 128
    print(randombytes(num_bytes).hex())
