"""HMAC-based PRNG with domain separation for PROTOCOL.

Uses HMAC-SHA256 to generate deterministic byte streams per domain tag.
This provides cryptographic separation between different random uses
while maintaining full determinism from (secret_key, public_seed).

Specification:
    K = UTF8(PROTOCOL_SECRET_KEY)
    S = UTF8(str(public_seed))
    D = UTF8(domain_tag)
    C = UTF8(str(counter))

    BLOCK(counter) = HMAC_SHA256(key=K, message= S || "|" || D || "|" || C)

The byte stream for a domain is BLOCK(0) || BLOCK(1) || ...
"""

from __future__ import annotations

import hmac
from typing import ClassVar


class DomainPRNG:
    """HMAC-SHA256 based PRNG with domain separation.

    Each domain maintains an independent byte stream with a (block_counter, offset)
    pointer. Bytes are consumed strictly in order within each domain.

    Example:
        >>> prng = DomainPRNG(secret_key="test-key", public_seed=12345)
        >>> prng.rand_u8("params")  # Get byte from "params" domain
        >>> prng.rand_int("map", 3, 5)  # Get int in [3, 5] from "map" domain
    """

    BLOCK_SIZE: ClassVar[int] = 32  # SHA256 output size

    def __init__(self, secret_key: str, public_seed: int) -> None:
        """Initialize PRNG from secret key and public seed.

        Args:
            secret_key: Secret key string (from PROTOCOL_SECRET_KEY env var)
            public_seed: Public seed integer (must be non-negative, <= 2^53-1)

        Raises:
            ValueError: If public_seed is invalid
        """
        if public_seed < 0:
            raise ValueError("public_seed must be non-negative")
        if public_seed > (2**53 - 1):
            raise ValueError("public_seed must be <= 2^53-1 for JSON portability")

        self._key = secret_key.encode()
        self._seed_str = str(public_seed)

        # Per-domain state: domain_tag -> (block_counter, offset, current_block)
        self._domains: dict[str, tuple[int, int, bytes]] = {}

    def _generate_block(self, domain: str, counter: int) -> bytes:
        """Generate a 32-byte block for a domain at a given counter.

        Args:
            domain: Domain tag string
            counter: Block counter (0, 1, 2, ...)

        Returns:
            32 bytes from HMAC-SHA256
        """
        # Message format: S || "|" || D || "|" || C
        # where S=seed, D=domain, C=counter (all as ASCII decimal, no leading zeros)
        message = f"{self._seed_str}|{domain}|{counter}".encode()
        return hmac.digest(self._key, message, "sha256")

    def _ensure_domain(self, domain: str) -> None:
        """Ensure domain is initialized with first block."""
        if domain not in self._domains:
            block = self._generate_block(domain, 0)
            self._domains[domain] = (0, 0, block)

    def _next_byte(self, domain: str) -> int:
        """Get next byte from domain stream.

        Args:
            domain: Domain tag string

        Returns:
            Next byte (0-255) from the domain's stream
        """
        self._ensure_domain(domain)
        block_counter, offset, block = self._domains[domain]

        byte_value = block[offset]
        offset += 1

        # Advance to next block if needed
        if offset >= self.BLOCK_SIZE:
            block_counter += 1
            offset = 0
            block = self._generate_block(domain, block_counter)

        self._domains[domain] = (block_counter, offset, block)
        return byte_value

    def rand_u8(self, domain: str) -> int:
        """Get next random byte (0-255) from domain stream.

        Args:
            domain: Domain tag string

        Returns:
            Random integer in [0, 255]
        """
        return self._next_byte(domain)

    def rand_int(self, domain: str, lo: int, hi: int) -> int:
        """Get unbiased random integer in [lo, hi] using rejection sampling.

        Uses 4 bytes (uint32) with rejection sampling to avoid modulo bias.

        Args:
            domain: Domain tag string
            lo: Lower bound (inclusive)
            hi: Upper bound (inclusive)

        Returns:
            Random integer in [lo, hi]

        Raises:
            ValueError: If lo > hi
        """
        if lo > hi:
            raise ValueError(f"lo ({lo}) must be <= hi ({hi})")

        n = hi - lo + 1
        if n == 1:
            return lo

        # Rejection sampling with uint32
        limit = (2**32 // n) * n  # Largest multiple of n that fits in uint32

        while True:
            # Read 4 bytes as big-endian uint32
            b0 = self._next_byte(domain)
            b1 = self._next_byte(domain)
            b2 = self._next_byte(domain)
            b3 = self._next_byte(domain)
            v = (b0 << 24) | (b1 << 16) | (b2 << 8) | b3

            if v < limit:
                return lo + (v % n)

    def peek_state(self, domain: str) -> tuple[int, int]:
        """Peek at current domain state (for debugging/testing).

        Args:
            domain: Domain tag string

        Returns:
            Tuple of (block_counter, offset)
        """
        self._ensure_domain(domain)
        block_counter, offset, _ = self._domains[domain]
        return (block_counter, offset)
