from __future__ import annotations

import math
from typing import Any

import pytest

import icsd.adapters.numba_invariants as numba_invariants
from icsd.adapters.numba_invariants import NumbaBatchFieldInvariant
from icsd.adapters.numpy_invariants import NumpyBatchFieldInvariant
from icsd.core.batch import validate_states
from icsd.core.invariants import Invariant, InvariantSet


def _assert_outcome_parity(
    expected: bool | tuple[Any, ...],
    actual: bool | tuple[Any, ...],
) -> None:
    if expected is True:
        assert actual is True
        return
    assert isinstance(expected, tuple)
    assert isinstance(actual, tuple)
    assert [_normalise_payload(item.as_dict()) for item in actual] == [
        _normalise_payload(item.as_dict()) for item in expected
    ]


def _normalise_payload(value: Any) -> Any:
    if isinstance(value, float) and math.isnan(value):
        return "NaN"
    if isinstance(value, list):
        return [_normalise_payload(item) for item in value]
    if isinstance(value, dict):
        return {key: _normalise_payload(item) for key, item in value.items()}
    return value


def _numpy_reference_helper(config: dict[str, Any]) -> NumpyBatchFieldInvariant:
    reference_config = dict(config)
    reference_config.setdefault("name", "numba_batch_field")
    reference_config.setdefault("failure_code", "numba_batch_invalid")
    return NumpyBatchFieldInvariant(**reference_config)


def test_constructor_rejects_invalid_expected_ndim_bool() -> None:
    with pytest.raises(TypeError, match="expected_ndim must be an integer"):
        NumbaBatchFieldInvariant(field="values", expected_ndim=True)


def test_non_mapping_failure_reports_stable_state_type_name() -> None:
    invariant = NumbaBatchFieldInvariant(field="values")

    outcome = invariant([1, 2, 3])  # type: ignore[arg-type]

    assert isinstance(outcome, tuple)
    assert len(outcome) == 1
    assert outcome[0].details == {
        "field": "values",
        "reason": "state_not_mapping",
        "state_type": "list",
    }


def test_missing_field_failure_envelope_is_deterministic() -> None:
    invariant = NumbaBatchFieldInvariant(field="values")

    outcome = invariant({"other": [1, 2, 3]})

    assert isinstance(outcome, tuple)
    assert len(outcome) == 1
    assert outcome[0].details == {
        "field": "values",
        "reason": "missing_field",
    }


def test_missing_numpy_returns_optional_dependency_failure(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(numba_invariants, "_load_numpy_runtime", lambda: None)
    invariant = NumbaBatchFieldInvariant(field="values")

    outcome = invariant({"values": [1, 2, 3]})

    assert isinstance(outcome, tuple)
    assert len(outcome) == 1
    failure = outcome[0]
    assert failure.code == "optional_dependency_missing"
    assert failure.details == {
        "dependency": "numpy",
        "extra": "speed-numpy",
        "install": "dpa-icsd[speed-numpy]",
        "field": "values",
        "reason": "optional_dependency_missing",
    }


def test_numba_missing_falls_back_to_numpy_reference(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    np = pytest.importorskip("numpy")
    numba_invariants._get_numba_kernels.cache_clear()
    monkeypatch.setattr(numba_invariants, "_load_numba_runtime", lambda: None)

    state = {"values": np.asarray([1.0, 2.0, 3.0], dtype=np.float64)}
    expected = _numpy_reference_helper({"field": "values"})(state)
    actual = NumbaBatchFieldInvariant(field="values")(state)

    _assert_outcome_parity(expected, actual)


def test_kernel_resolution_exception_falls_back_without_raising(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    np = pytest.importorskip("numpy")

    def _raise_kernel_error() -> Any:
        raise RuntimeError("kernel resolution failure")

    monkeypatch.setattr(numba_invariants, "_get_numba_kernels", _raise_kernel_error)

    state = {"values": np.asarray([-2, 1], dtype=np.int64)}
    expected = _numpy_reference_helper({"field": "values", "min_value": 0})(state)
    actual = NumbaBatchFieldInvariant(field="values", min_value=0)(state)

    _assert_outcome_parity(expected, actual)


def test_parity_with_numpy_helper_for_reference_scenarios() -> None:
    np = pytest.importorskip("numpy")
    scenarios: list[tuple[dict[str, Any], dict[str, Any]]] = [
        (
            {
                "field": "values",
                "expected_ndim": 1,
                "min_value": 0,
                "max_value": 10,
                "require_finite": True,
            },
            {"values": np.asarray([1.0, 2.0, 3.0], dtype=np.float64)},
        ),
        (
            {"field": "values"},
            {"values": np.asarray([True, False], dtype=np.bool_)},
        ),
        (
            {"field": "values", "expected_ndim": 1},
            {"values": np.asarray([[1.0, 2.0]], dtype=np.float64)},
        ),
        (
            {"field": "values", "allow_empty": False},
            {"values": np.asarray([], dtype=np.float64)},
        ),
        (
            {"field": "values", "require_finite": True, "min_value": 0},
            {"values": np.asarray([float("nan"), 1.0], dtype=np.float64)},
        ),
        (
            {"field": "values", "min_value": 0},
            {"values": np.asarray([-3.0, 5.0], dtype=np.float64)},
        ),
        (
            {"field": "values", "max_value": 2},
            {"values": np.asarray([1.0, 4.0], dtype=np.float64)},
        ),
        (
            {"field": "values", "require_finite": False, "min_value": 0},
            {"values": np.asarray([float("nan"), -1.0], dtype=np.float64)},
        ),
    ]

    for config, state in scenarios:
        expected = _numpy_reference_helper(config)(state)
        actual = NumbaBatchFieldInvariant(**config)(state)
        _assert_outcome_parity(expected, actual)


def test_nan_with_require_finite_false_uses_reference_path_before_numba_bounds(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    np = pytest.importorskip("numpy")

    def _raise_if_called(*_: Any, **__: Any) -> Any:
        raise AssertionError("Numba min/max path must not run for float NaN bounds case.")

    monkeypatch.setattr(numba_invariants, "_scan_min_max_with_optional_numba", _raise_if_called)

    config = {
        "field": "values",
        "require_finite": False,
        "min_value": 0,
    }
    state = {"values": np.asarray([float("nan"), -1.0], dtype=np.float64)}

    expected = _numpy_reference_helper(config)(state)
    actual = NumbaBatchFieldInvariant(**config)(state)

    _assert_outcome_parity(expected, actual)


def test_numba_path_opportunistic_parity() -> None:
    np = pytest.importorskip("numpy")
    pytest.importorskip("numba")
    numba_invariants._get_numba_kernels.cache_clear()

    config = {
        "field": "values",
        "expected_ndim": 1,
        "min_value": 0,
        "max_value": 10,
        "require_finite": True,
    }
    state = {"values": np.asarray([1.0, 2.0, 3.0], dtype=np.float64)}

    expected = _numpy_reference_helper(config)(state)
    first = NumbaBatchFieldInvariant(**config)(state)
    second = NumbaBatchFieldInvariant(**config)(state)

    _assert_outcome_parity(expected, first)
    _assert_outcome_parity(first, second)


def test_batch_integration_matches_numpy_reference_report() -> None:
    np = pytest.importorskip("numpy")

    config = {
        "field": "values",
        "expected_ndim": 1,
        "min_value": 0,
        "require_finite": True,
    }
    states = [
        {"values": np.asarray([1.0, 2.0], dtype=np.float64)},
        {"values": np.asarray([float("nan")], dtype=np.float64)},
        {"values": np.asarray([-3.0], dtype=np.float64)},
        {"other": [1, 2, 3]},
    ]

    numpy_set = InvariantSet(
        [Invariant(name="values_check", check=_numpy_reference_helper(config))]
    )
    numba_set = InvariantSet(
        [Invariant(name="values_check", check=NumbaBatchFieldInvariant(**config))]
    )

    numpy_report = validate_states(numpy_set, states)
    numba_report = validate_states(numba_set, states)

    assert [result.index for result in numba_report.results] == [0, 1, 2, 3]
    assert numba_report.as_dict() == numpy_report.as_dict()
