"""Deterministic random-state management.

The generator must produce byte-identical output for a given ``(master_seed,
sample_index)`` pair regardless of process/thread scheduling, batch size, or
worker count. To guarantee this we never share a single global RNG across
samples or pipeline stages. Instead, every stage of every sample draws from
an independently *derived* RNG whose seed is a pure function of
``(master_seed, sample_index, stage_name)``.

This makes the pipeline safe for multiprocessing: worker N generating sample
57 produces exactly the same random draws as a single-process run would have,
because nothing depends on execution order.
"""

from __future__ import annotations

import hashlib
import random
from typing import Any

import numpy as np

_MOD32 = 2**32


def derive_seed(master_seed: int, *parts: Any) -> int:
    """Derive a deterministic 32-bit seed from a master seed and arbitrary parts.

    Uses SHA-256 over the concatenated string representation of all parts so
    that the mapping is stable across Python versions and processes (unlike
    Python's salted ``hash()``).

    Args:
        master_seed: The top-level seed configured by the user.
        *parts: Any number of hashable/stringifiable disambiguators, e.g.
            sample index and pipeline stage name.

    Returns:
        A non-negative integer in ``[0, 2**32)`` suitable for seeding
        ``numpy.random.Generator`` or ``random.Random``.
    """
    key = f"{master_seed}|" + "|".join(str(p) for p in parts)
    digest = hashlib.sha256(key.encode("utf-8")).digest()
    return int.from_bytes(digest[:4], byteorder="big") % _MOD32


class SeedManager:
    """Provides independent, deterministic RNGs for each pipeline stage of one sample.

    Example:
        >>> sm = SeedManager(master_seed=42, sample_index=7)
        >>> rng = sm.numpy_rng("placement")
        >>> rng.uniform(0, 1)
    """

    def __init__(self, master_seed: int, sample_index: int) -> None:
        self.master_seed = master_seed
        self.sample_index = sample_index

    def seed_for(self, stage: str) -> int:
        """Return the deterministic integer seed for a named pipeline stage."""
        return derive_seed(self.master_seed, self.sample_index, stage)

    def numpy_rng(self, stage: str) -> np.random.Generator:
        """Return a fresh ``numpy.random.Generator`` (PCG64) seeded for ``stage``."""
        return np.random.default_rng(self.seed_for(stage))

    def python_rng(self, stage: str) -> random.Random:
        """Return a fresh ``random.Random`` seeded for ``stage``."""
        return random.Random(self.seed_for(stage))


def resolve_scalar_or_range(value: Any, rng: np.random.Generator) -> float:
    """Resolve a config parameter that is either a fixed scalar or a ``{min, max}`` range.

    Args:
        value: Either a plain number, or an object/mapping exposing ``min``
            and ``max`` attributes/keys (a :class:`~src.config.schema.Range`
            or an equivalent dict).
        rng: The seeded generator to draw from when ``value`` is a range.

    Returns:
        A resolved float value.
    """
    if isinstance(value, (int, float)):
        return float(value)
    if isinstance(value, dict) and "min" in value and "max" in value:
        lo, hi = float(value["min"]), float(value["max"])
        return float(rng.uniform(lo, hi)) if hi > lo else lo
    # Pydantic Range model (duck-typed access)
    lo = getattr(value, "min", None)
    hi = getattr(value, "max", None)
    if lo is not None and hi is not None:
        lo, hi = float(lo), float(hi)
        return float(rng.uniform(lo, hi)) if hi > lo else lo
    raise TypeError(f"Cannot resolve scalar-or-range value: {value!r}")