"""Optional Numba-backed invariant helper with deterministic NumPy fallback.

This helper follows a strict first-failure-wins contract: each call returns either
True or exactly one InvariantFailure. Numba is best-effort only; missing,
unsupported, or failing Numba paths always fall back to NumPy reference checks.
"""

from __future__ import annotations

import math
from collections.abc import Mapping
from dataclasses import dataclass
from functools import lru_cache
from typing import Any

from icsd.core.errors import InvariantFailure

from .numpy_invariants import ACCEPTED_DTYPE_KINDS, NumpyBatchFieldInvariant

__all__ = ["NumbaBatchFieldInvariant"]


@dataclass(frozen=True, slots=True)
class NumbaBatchFieldInvariant:
    """Declarative Numba experiment helper with deterministic reference fallback."""

    field: str
    name: str = "numba_batch_field"
    failure_code: str = "numba_batch_invalid"
    expected_ndim: int | None = None
    min_value: int | float | None = None
    max_value: int | float | None = None
    require_finite: bool = True
    allow_empty: bool = True

    def __post_init__(self) -> None:
        # Reuse WS-0077 constructor contract validation.
        self._reference_helper()

    def __call__(self, state: Mapping[str, Any]) -> bool | tuple[InvariantFailure, ...]:
        reference = self._reference_helper()

        if not isinstance(state, Mapping):
            return reference(state)
        if self.field not in state:
            return reference(state)

        numpy_runtime = _load_numpy_runtime()
        if numpy_runtime is None:
            return (
                InvariantFailure(
                    invariant=self.name,
                    code="optional_dependency_missing",
                    message="Optional dependency 'numpy' is required for this invariant.",
                    details={
                        "dependency": "numpy",
                        "extra": "speed-numpy",
                        "install": "dpa-icsd[speed-numpy]",
                        "field": self.field,
                        "reason": "optional_dependency_missing",
                    },
                ),
            )

        try:
            array = numpy_runtime.asarray(state[self.field])
        except Exception:
            return reference(state)

        dtype_kind = str(array.dtype.kind)
        if dtype_kind not in ACCEPTED_DTYPE_KINDS:
            return reference(state)

        if self.expected_ndim is not None and int(array.ndim) != self.expected_ndim:
            return reference(state)

        is_empty = int(array.size) == 0
        if is_empty:
            if not self.allow_empty:
                return reference(state)
            return True

        if self.require_finite and not _all_finite_with_optional_numba(
            array,
            np_runtime=numpy_runtime,
        ):
            return reference(state)

        has_bounds = self.min_value is not None or self.max_value is not None

        # For require_finite=False on float arrays, preserve reference bound semantics
        # by using the reference path whenever NaNs are present.
        if (
            has_bounds
            and not self.require_finite
            and dtype_kind == "f"
            and bool(numpy_runtime.isnan(array).any())
        ):
            return reference(state)

        if has_bounds:
            observed_min, observed_max = _scan_min_max_with_optional_numba(
                array,
                np_runtime=numpy_runtime,
            )
            if self.min_value is not None and observed_min < self.min_value:
                return reference(state)
            if self.max_value is not None and observed_max > self.max_value:
                return reference(state)

        return True

    def _reference_helper(self) -> NumpyBatchFieldInvariant:
        return NumpyBatchFieldInvariant(
            field=self.field,
            name=self.name,
            failure_code=self.failure_code,
            expected_ndim=self.expected_ndim,
            min_value=self.min_value,
            max_value=self.max_value,
            require_finite=self.require_finite,
            allow_empty=self.allow_empty,
        )


def _load_numpy_runtime() -> Any | None:
    try:
        import numpy as np
    except Exception:
        return None
    return np


def _load_numba_runtime() -> Any | None:
    try:
        import numba
    except Exception:
        return None
    return numba


@lru_cache(maxsize=1)
def _get_numba_kernels() -> tuple[Any, Any] | None:
    numba_runtime = _load_numba_runtime()
    if numba_runtime is None:
        return None

    try:

        @numba_runtime.njit(cache=True)
        def _all_finite_kernel(values: Any) -> bool:
            for index in range(values.size):
                if not math.isfinite(values[index]):
                    return False
            return True

        @numba_runtime.njit(cache=True)
        def _min_max_kernel(values: Any) -> tuple[Any, Any]:
            minimum = values[0]
            maximum = values[0]
            for index in range(1, values.size):
                value = values[index]
                if value < minimum:
                    minimum = value
                if value > maximum:
                    maximum = value
            return minimum, maximum

    except Exception:
        return None

    return _all_finite_kernel, _min_max_kernel


def _all_finite_with_optional_numba(array: Any, *, np_runtime: Any) -> bool:
    # Kernels are only used after dtype/ndim/empty gates pass.
    flat = array.reshape(-1)

    try:
        kernels = _get_numba_kernels()
        if kernels is not None:
            return bool(kernels[0](flat))
    except Exception:
        pass

    return bool(np_runtime.isfinite(array).all())


def _scan_min_max_with_optional_numba(
    array: Any,
    *,
    np_runtime: Any,
) -> tuple[int | float, int | float]:
    # Kernels are only used after dtype/ndim/empty gates pass.
    flat = array.reshape(-1)

    try:
        kernels = _get_numba_kernels()
        if kernels is not None:
            minimum, maximum = kernels[1](flat)
            return (
                _to_py_scalar(minimum, np_runtime=np_runtime),
                _to_py_scalar(maximum, np_runtime=np_runtime),
            )
    except Exception:
        pass

    return (
        _to_py_scalar(array.min(), np_runtime=np_runtime),
        _to_py_scalar(array.max(), np_runtime=np_runtime),
    )


def _to_py_scalar(value: Any, *, np_runtime: Any) -> int | float:
    raw_value = value.item() if hasattr(value, "item") else value
    if isinstance(raw_value, np_runtime.integer) or (
        isinstance(raw_value, int) and not isinstance(raw_value, bool)
    ):
        return int(raw_value)
    if isinstance(raw_value, np_runtime.floating) or isinstance(raw_value, float):
        return float(raw_value)
    if isinstance(raw_value, bool):
        return int(raw_value)
    msg = "Expected a numeric scalar value."
    raise TypeError(msg)
