from __future__ import annotations

"""
rng.py — OS RNG wrapper & deterministic test streams (Animica PQ utils)

Goals
-----
- Cryptographically secure randomness via the OS (secrets/os.urandom).
- Deterministic, reproducible byte streams for tests/fixtures/simulations.
- Small, dependency-free, Python-standard-library only.
- Ergonomic helpers for common needs (u64, uniform range, hex seeds).

Design
------
Secure random:
  * `secure_random(n)` -> bytes
  * `os_random(n)` -> bytes (alias)

Deterministic stream:
  * `DeterministicStream(seed: bytes, label: bytes=b"rng-v1")`
    - Expand using SHAKE-256 with a domain tag and a 64-bit counter.
    - Not a KDF substitute; only for test determinism & simulations.
  * `deterministic_rng(seed: BytesLike, *, label: BytesLike=b"rng-v1") -> DeterministicStream`

Helpers:
  * `random_u64(stream=None)`           -> 64-bit uint (secure if stream=None)
  * `random_uniform(n, stream=None)`    -> unbiased int in [0, n) (n>0)
  * `derive_seed_from_hex(s)`           -> bytes, even/odd hex len supported

Security notes
--------------
- Never use DeterministicStream for secrets in production. It is *not* a CSPRNG.
- For cryptographic keys, nonces, salts, etc., always use `secure_random`.
"""

import hashlib
import os
import secrets
import struct
import threading
from typing import Optional, Union

BytesLike = Union[bytes, bytearray, memoryview, str]


# ---------------------------------------------------------------------------
# Secure RNG
# ---------------------------------------------------------------------------


def secure_random(n: int) -> bytes:
    """
    Return `n` bytes from the OS CSPRNG.
    """
    if n < 0:
        raise ValueError("n must be non-negative")
    # secrets.token_bytes may use os.urandom under the hood; it's the highest-level API.
    return secrets.token_bytes(n)


def os_random(n: int) -> bytes:
    """
    Alias of `secure_random` for clarity.
    """
    return secure_random(n)


# ---------------------------------------------------------------------------
# Deterministic test RNG (SHAKE-256 stream)
# ---------------------------------------------------------------------------

_DOMAIN = b"animica:pq:utils:rng:shake256:v1"


def _to_bytes(b: BytesLike, *, encoding: str = "utf-8") -> bytes:
    """
    Accept bytes/bytearray/memoryview or hex-string (0x...) / ascii and return bytes.
    - If `b` is a str that looks like hex (with or without 0x), decode as hex.
    - Else encode as UTF-8.
    """
    if isinstance(b, (bytes, bytearray)):
        return bytes(b)
    if isinstance(b, memoryview):
        return b.tobytes()
    if isinstance(b, str):
        s = b.strip()
        # Hex-ish? 0x... or [0-9a-fA-F]+ with even/odd len accepted.
        if s.startswith("0x") or all(ch in "0123456789abcdefABCDEF" for ch in s):
            if s.startswith("0x"):
                s = s[2:]
            if len(s) % 2 == 1:
                s = "0" + s
            try:
                return bytes.fromhex(s)
            except ValueError:
                # Not valid hex; fall back to UTF-8
                pass
        return s.encode(encoding)
    raise TypeError(f"unsupported seed type: {type(b)!r}")


class DeterministicStream:
    """
    An infinite deterministic byte stream based on SHAKE-256.

    Construction:
        rng = DeterministicStream(seed=b"...", label=b"rng-v1")

    API:
        rng.read(n)        -> bytes
        rng.random_u64()   -> int
        rng.random_uniform(n) -> int in [0, n)

    Internals:
        Bytes are generated by SHAKE256(domain || label_len || label || seed_len || seed || counter_be64)
        in 1024-byte blocks. `counter` increments per block.
    """

    __slots__ = ("_seed", "_label", "_counter", "_buf", "_lock")

    _BLOCK = 1024  # refill size

    def __init__(self, seed: BytesLike, label: BytesLike = b"rng-v1"):
        seed_b = _to_bytes(seed)
        label_b = _to_bytes(label)
        if len(seed_b) == 0:
            raise ValueError("DeterministicStream seed must be non-empty")
        self._seed = seed_b
        self._label = label_b
        self._counter = 0
        self._buf = bytearray()
        self._lock = threading.Lock()

    def _refill(self) -> None:
        self._counter += 1
        # Frame transcript for this block
        #   D || len(label) || label || len(seed) || seed || counter_be64
        frame = (
            _DOMAIN
            + struct.pack(">I", len(self._label))
            + self._label
            + struct.pack(">I", len(self._seed))
            + self._seed
            + struct.pack(">Q", self._counter)
        )
        xof = hashlib.shake_256(frame)
        self._buf.extend(xof.digest(self._BLOCK))

    def read(self, n: int) -> bytes:
        if n < 0:
            raise ValueError("n must be non-negative")
        with self._lock:
            while len(self._buf) < n:
                self._refill()
            out = bytes(self._buf[:n])
            del self._buf[:n]
            return out

    def random_u64(self) -> int:
        """
        Return a 64-bit unsigned integer from the stream.
        """
        return struct.unpack(">Q", self.read(8))[0]

    def random_uniform(self, n: int) -> int:
        """
        Unbiased integer in [0, n). Uses rejection sampling.
        """
        if n <= 0:
            raise ValueError("n must be positive")
        # Compute rejection bound based on 64-bit space.
        limit = (1 << 64) - ((1 << 64) % n)
        while True:
            x = self.random_u64()
            if x < limit:
                return x % n


def deterministic_rng(
    seed: BytesLike, *, label: BytesLike = b"rng-v1"
) -> DeterministicStream:
    """
    Convenience constructor.
    """
    return DeterministicStream(seed, label)


# ---------------------------------------------------------------------------
# Helper functions (secure or deterministic)
# ---------------------------------------------------------------------------


def random_u64(stream: Optional[DeterministicStream] = None) -> int:
    """
    If `stream` provided, use it; else draw 8 bytes from OS RNG.
    """
    if stream is None:
        return struct.unpack(">Q", secure_random(8))[0]
    return stream.random_u64()


def random_uniform(n: int, stream: Optional[DeterministicStream] = None) -> int:
    """
    Unbiased random int in [0, n).
    """
    if stream is None:
        # secure path
        if n <= 0:
            raise ValueError("n must be positive")
        # rejection sampling with OS RNG
        limit = (1 << 64) - ((1 << 64) % n)
        while True:
            x = struct.unpack(">Q", secure_random(8))[0]
            if x < limit:
                return x % n
    else:
        return stream.random_uniform(n)


def derive_seed_from_hex(s: str) -> bytes:
    """
    Convenience to interpret a string as hex (with or without 0x).
    Falls back to utf-8 if not hex-like.
    """
    return _to_bytes(s)


# ---------------------------------------------------------------------------
# Self-test / demonstration
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    print("[rng] Secure 16 bytes:", os.urandom(16).hex())

    # Deterministic demo
    seed = "0xdeadbeefcafebabe"
    rng1 = deterministic_rng(seed, label=b"demo")
    rng2 = deterministic_rng(seed, label=b"demo")
    a = rng1.read(32).hex()
    b = rng2.read(16).hex() + rng2.read(16).hex()  # split reads should match
    print("[rng] Deterministic equality:", "OK" if a == b else "MISMATCH")
    print("[rng] random_u64 example:", hex(rng1.random_u64()))
    print("[rng] uniform[0,10):", [rng1.random_uniform(10) for _ in range(8)])
