"""Order-independent named NumPy ``SeedSequence`` substreams."""

from __future__ import annotations

import hashlib
from typing import Iterable, Mapping

import numpy as np

from .hashing import canonical_json_bytes


DEFAULT_STREAMS = ("trajectory", "sensor", "model", "network")


def _label_words(*labels: object) -> list[int]:
    digest = hashlib.sha256(canonical_json_bytes(labels)).digest()
    return [
        int.from_bytes(digest[offset : offset + 4], "little")
        for offset in range(0, 16, 4)
    ]


def named_seed_record(
    root_seed: int,
    namespace: object,
    stream_names: Iterable[str] = DEFAULT_STREAMS,
) -> dict[str, list[int]]:
    """Create stable named seed states independent of request order."""
    root_seed = int(root_seed)
    if root_seed < 0:
        raise ValueError("root_seed must be non-negative")
    record: dict[str, list[int]] = {}
    for name in sorted(set(stream_names)):
        if not name:
            raise ValueError("stream names must be non-empty")
        entropy = [root_seed & 0xFFFFFFFF, (root_seed >> 32) & 0xFFFFFFFF]
        entropy.extend(_label_words(namespace, name))
        state = np.random.SeedSequence(entropy).generate_state(4, dtype=np.uint32)
        record[name] = [int(word) for word in state]
    return record


def generator_from_record(
    record: Mapping[str, list[int]],
    stream_name: str,
) -> np.random.Generator:
    if stream_name not in record:
        raise KeyError(f"Unknown random stream {stream_name!r}")
    words = [int(word) for word in record[stream_name]]
    if len(words) != 4 or any(word < 0 or word > 0xFFFFFFFF for word in words):
        raise ValueError(f"Invalid seed state for stream {stream_name!r}")
    return np.random.default_rng(np.random.SeedSequence(words))