"""
RGE-256 Core PRNG Implementation

256-bit ARX-based pseudorandom number generator with geometric entropy
derived from Recursive Division Tree (RDT) analysis.
"""

import hashlib
from typing import Tuple, List, Optional, Union


def rotl32(x: int, r: int) -> int:
    """32-bit left rotation."""
    x = x & 0xFFFFFFFF
    r = r % 32
    return ((x << r) | (x >> (32 - r))) & 0xFFFFFFFF


def simple_hash(s: str) -> int:
    """Simple string hash for domain separation (FNV-1a)."""
    h = 0x811c9dc5
    for c in s:
        h ^= ord(c)
        h = (h * 0x01000193) & 0xFFFFFFFF
    return h


class RGE256:
    """
    RGE-256 Pseudorandom Number Generator.
    
    A novel ARX-style PRNG using rotation schedules derived from
    Recursive Division Tree entropy constants (zetas).
    """
    
    DEFAULT_ZETAS = (1.585, 1.926, 1.262)
    K_BASE = [
        0x9E3779B9, 0x517CC1B7, 0xC2B2AE35, 0x165667B1,
        0x85EBCA77, 0x27D4EB2F, 0xDE5FB9D7, 0x94D049BB
    ]
    
    def __init__(
        self,
        seed: Optional[Union[int, bytes, str]] = None,
        rounds: int = 3,
        zetas: Tuple[float, float, float] = None,
        domain: str = "rge256-default"
    ):
        self.rounds = max(1, min(10, rounds))
        self.zetas = zetas if zetas is not None else self.DEFAULT_ZETAS
        self.domain = domain
        
        self.s0 = self.s1 = self.s2 = self.s3 = 0
        self.s4 = self.s5 = self.s6 = self.s7 = 0
        self.r = [0] * 8
        self.k = [0] * 8
        
        self._seed(seed)
    
    def _seed(self, seed: Optional[Union[int, bytes, str]] = None) -> None:
        if seed is None:
            import os
            seed_bytes = os.urandom(32)
        elif isinstance(seed, int):
            seed_bytes = seed.to_bytes(8, 'big', signed=False)
        elif isinstance(seed, str):
            seed_bytes = seed.encode('utf-8')
        elif isinstance(seed, bytes):
            seed_bytes = seed
        else:
            seed_bytes = str(seed).encode('utf-8')
        
        seed_bytes = self.domain.encode('utf-8') + seed_bytes
        h = hashlib.sha512(seed_bytes).digest()
        
        state = [int.from_bytes(h[i*4:(i+1)*4], 'little') for i in range(16)]
        
        self.s0, self.s1, self.s2, self.s3 = state[0], state[1], state[2], state[3]
        self.s4, self.s5, self.s6, self.s7 = state[4], state[5], state[6], state[7]
        
        tri, meng, tet = self.zetas
        base = [
            tri, meng, tet, (tri + meng + tet) / 3.0,
            tri + 0.5, meng + 0.75, tet + 0.25, (tri * 1.25 + tet * 0.75) / 2
        ]
        self.r = [max(1, int(abs(z * 977 + (i * 7 + 13)) % 31) or 1) for i, z in enumerate(base)]
        self.k = [(self.K_BASE[i] ^ state[8 + i]) & 0xFFFFFFFF for i in range(8)]
        
        for _ in range(10):
            self._step()
    
    def _step(self) -> None:
        s0, s1, s2, s3 = self.s0, self.s1, self.s2, self.s3
        s4, s5, s6, s7 = self.s4, self.s5, self.s6, self.s7
        r0, r1, r2, r3, r4, r5, r6, r7 = self.r
        
        # Quad A
        s0 = (s0 + s1 + self.k[0]) & 0xFFFFFFFF
        s1 = rotl32(s1 ^ s0, r0)
        s2 = (s2 + s3 + self.k[1]) & 0xFFFFFFFF
        s3 = rotl32(s3 ^ s2, r1)
        tA0, tA1 = s0 ^ s2, s1 ^ s3
        s0 = (s0 + rotl32(tA0, r2) + self.k[2]) & 0xFFFFFFFF
        s2 = (s2 + rotl32(tA1, r3) + self.k[3]) & 0xFFFFFFFF
        s1 ^= rotl32(s0, (r0 + r2) % 31 or 1)
        s3 ^= rotl32(s2, (r1 + r3) % 31 or 1)
        
        # Quad B
        s4 = (s4 + s5 + self.k[4]) & 0xFFFFFFFF
        s5 = rotl32(s5 ^ s4, r4)
        s6 = (s6 + s7 + self.k[5]) & 0xFFFFFFFF
        s7 = rotl32(s7 ^ s6, r5)
        tB0, tB1 = s4 ^ s6, s5 ^ s7
        s4 = (s4 + rotl32(tB0, r6) + self.k[6]) & 0xFFFFFFFF
        s6 = (s6 + rotl32(tB1, r7) + self.k[7]) & 0xFFFFFFFF
        s5 ^= rotl32(s4, (r4 + r6) % 31 or 1)
        s7 ^= rotl32(s6, (r5 + r7) % 31 or 1)
        
        # Cross-quad mixing
        s1 ^= rotl32(s5, 13)
        s3 ^= rotl32(s7, 7)
        s5 ^= rotl32(s1, 11)
        s7 ^= rotl32(s3, 17)
        
        # Global fold
        m0, m1, m2, m3 = s0 ^ s4, s1 ^ s5, s2 ^ s6, s3 ^ s7
        s0 = (s0 + rotl32(m1, 3)) & 0xFFFFFFFF
        s1 = (s1 + rotl32(m2, 5)) & 0xFFFFFFFF
        s2 = (s2 + rotl32(m3, 7)) & 0xFFFFFFFF
        s3 = (s3 + rotl32(m0, 11)) & 0xFFFFFFFF
        s4 ^= rotl32(s0, 19)
        s5 ^= rotl32(s1, 23)
        s6 ^= rotl32(s2, 29)
        s7 ^= rotl32(s3, 31)
        
        self.s0, self.s1, self.s2, self.s3 = s0, s1, s2, s3
        self.s4, self.s5, self.s6, self.s7 = s4, s5, s6, s7
    
    def next32(self) -> int:
        for _ in range(self.rounds):
            self._step()
        return (self.s0 ^ rotl32(self.s4, 13)) & 0xFFFFFFFF
    
    def next64(self) -> int:
        return (self.next32() << 32) | self.next32()
    
    def next_float(self) -> float:
        return self.next32() / 0x100000000
    
    def next_double(self) -> float:
        return self.next64() / (2**64)
    
    def next_range(self, min_val: int, max_val: int) -> int:
        if min_val > max_val:
            min_val, max_val = max_val, min_val
        return min_val + (self.next32() % (max_val - min_val + 1))
    
    def next_bytes(self, n: int) -> bytes:
        result = bytearray()
        while len(result) < n:
            result.extend(self.next32().to_bytes(4, 'little'))
        return bytes(result[:n])
    
    def shuffle(self, sequence: list) -> list:
        result = list(sequence)
        for i in range(len(result) - 1, 0, -1):
            j = self.next_range(0, i)
            result[i], result[j] = result[j], result[i]
        return result
    
    def choice(self, sequence: list):
        if not sequence:
            raise ValueError("Cannot choose from empty sequence")
        return sequence[self.next_range(0, len(sequence) - 1)]
    
    def sample(self, sequence: list, k: int) -> list:
        if k > len(sequence):
            raise ValueError("Sample size larger than population")
        return self.shuffle(sequence)[:k]
    
    def get_state(self) -> dict:
        return {
            'state': [self.s0, self.s1, self.s2, self.s3, self.s4, self.s5, self.s6, self.s7],
            'rotations': list(self.r),
            'keys': list(self.k),
            'rounds': self.rounds,
            'zetas': self.zetas,
            'domain': self.domain,
            'version': '1.0.7'
        }
    
    def set_state(self, state_dict: dict) -> None:
        s = state_dict['state']
        self.s0, self.s1, self.s2, self.s3 = s[0], s[1], s[2], s[3]
        self.s4, self.s5, self.s6, self.s7 = s[4], s[5], s[6], s[7]
        self.r = list(state_dict['rotations'])
        self.k = list(state_dict['keys'])
        self.rounds = state_dict['rounds']
        self.zetas = tuple(state_dict['zetas'])
        self.domain = state_dict['domain']
    
    def __repr__(self) -> str:
        return f"RGE256(rounds={self.rounds}, zetas={self.zetas}, domain='{self.domain}')"
