#!/usr/bin/env python
"""A module with utilities for optimized pseudo-random number generation."""

import os
import struct
import threading
from typing import Callable, List

_random_buffer_size: int = 1024
_random_buffer: List[int] = []
_mutex = threading.Lock()


def UInt16() -> int:
  """Returns a pseudo-random 16-bit unsigned integer."""
  return UInt32() & 0xFFFF


def PositiveUInt16() -> int:
  """Returns a pseudo-random 16-bit non-zero unsigned integer."""
  return _Positive(UInt16)


def UInt32() -> int:
  """Returns a pseudo-random 32-bit unsigned integer."""
  with _mutex:
    try:
      return _random_buffer.pop()
    except IndexError:
      data = os.urandom(struct.calcsize("=L") * _random_buffer_size)
      _random_buffer.extend(
          struct.unpack("=" + "L" * _random_buffer_size, data)
      )
      return _random_buffer.pop()


def PositiveUInt32() -> int:
  """Returns a pseudo-random 32-bit non-zero unsigned integer."""
  return _Positive(UInt32)


def UInt64() -> int:
  """Returns a pseudo-random 64-bit unsigned integer."""
  return (UInt32() << 32) | UInt32()


def Id64() -> int:
  """Returns a pseudo-random, positive, non-zero 64-bit int usable as ID.

  The returned int is guaranteed to have at least one of the most significant 32
  bits set to 1. Effectively, all IDs generated by this function are
  >= 0x100000000. This is due to backwards compatibility: Legacy, 32-bit IDs can
  be identified by having 0 in the 32 most significant bits. Legacy IDs are
  serialized as 8-character hexadecimal, whereas new IDs are serialized as 16-
  character hexadecimal string.
  """
  return (PositiveUInt32() << 32) | UInt32()


def _Positive(rng: Callable[[], int]) -> int:
  while True:
    result = rng()
    if result > 0:
      return result