"""Tests for :mod:`dlm_sway.core.determinism`."""

from __future__ import annotations

import os
import random

import numpy as np

from dlm_sway.core.determinism import DeterminismSummary, seed_everything


class TestSeedEverything:
    def test_returns_summary(self) -> None:
        summary = seed_everything(0)
        assert isinstance(summary, DeterminismSummary)
        assert summary.seed == 0
        assert summary.class_ in {"strict", "best_effort", "loose"}

    def test_idempotent_for_stdlib_random(self) -> None:
        seed_everything(42)
        a = [random.random() for _ in range(5)]
        seed_everything(42)
        b = [random.random() for _ in range(5)]
        assert a == b

    def test_idempotent_for_numpy(self) -> None:
        seed_everything(17)
        a = np.random.rand(5)
        seed_everything(17)
        b = np.random.rand(5)
        np.testing.assert_array_equal(a, b)

    def test_cublas_workspace_set_under_strict(self) -> None:
        os.environ.pop("CUBLAS_WORKSPACE_CONFIG", None)
        seed_everything(0, strict=True)
        assert os.environ.get("CUBLAS_WORKSPACE_CONFIG") == ":4096:8"

    def test_non_strict_does_not_set_cublas(self) -> None:
        os.environ.pop("CUBLAS_WORKSPACE_CONFIG", None)
        seed_everything(0, strict=False)
        # Non-strict mode must not leak the env var in either direction;
        # the host environment's prior value wins.
        assert (
            "CUBLAS_WORKSPACE_CONFIG" not in os.environ
            or os.environ["CUBLAS_WORKSPACE_CONFIG"] != ":4096:8"
        )


class TestRunnerCallsSeedEverything:
    """The runner must seed every RNG before any probe runs (P09)."""

    def test_runner_populates_determinism_field(self) -> None:
        from dlm_sway.backends.dummy import DummyDifferentialBackend, DummyResponses
        from dlm_sway.suite.runner import run as run_suite
        from dlm_sway.suite.spec import SwaySpec

        backend = DummyDifferentialBackend(base=DummyResponses(), ft=DummyResponses())
        spec = SwaySpec.model_validate(
            {
                "version": 1,
                "models": {
                    "base": {"base": "b"},
                    "ft": {"base": "b", "adapter": "/tmp/a"},
                },
                "defaults": {"seed": 7},
                "suite": [],
            }
        )
        result = run_suite(spec, backend)
        assert result.determinism is not None
        assert result.determinism.seed == 7
        assert result.determinism.class_ in {"strict", "best_effort", "loose"}

    def test_runner_seeds_before_first_probe(self, monkeypatch) -> None:
        """Reorder check: seed_everything must fire *before* the probe loop."""
        from dlm_sway.backends.dummy import DummyDifferentialBackend, DummyResponses
        from dlm_sway.suite import runner as runner_mod
        from dlm_sway.suite.spec import SwaySpec

        events: list[str] = []

        original_seed = runner_mod.seed_everything

        def recording_seed(seed: int, *, strict: bool = True):
            events.append(f"seed={seed}")
            return original_seed(seed, strict=strict)

        monkeypatch.setattr(runner_mod, "seed_everything", recording_seed)

        # Use the dummy preflight as a probe stand-in: it runs *after*
        # seeding in the runner, so its event lands after the seed event.
        from dlm_sway.backends import dummy as dummy_mod

        original_preflight = dummy_mod.DummyDifferentialBackend.preflight_finite_check

        def recording_preflight(self):
            events.append("preflight")
            return original_preflight(self)

        monkeypatch.setattr(
            dummy_mod.DummyDifferentialBackend,
            "preflight_finite_check",
            recording_preflight,
        )

        backend = DummyDifferentialBackend(base=DummyResponses(), ft=DummyResponses())
        spec = SwaySpec.model_validate(
            {
                "version": 1,
                "models": {
                    "base": {"base": "b"},
                    "ft": {"base": "b", "adapter": "/tmp/a"},
                },
                "defaults": {"seed": 11},
                "suite": [],
            }
        )
        runner_mod.run(spec, backend)
        assert events.index("seed=11") < events.index("preflight"), (
            f"seed must fire before preflight; got {events}"
        )