"""Internal mask implementations and scalar-to-mask coercion helpers.

The public mask API lives mostly on :class:`topovec.ansatz.base.MaskND`. This
module provides the concrete private helpers that back callable masks and
automatic scalar coercion in mask algebra and ansatz masking.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import numpy as np

from .base import MaskND
from .frame import CoordinateSampleND, CoordinateSystemND

__all__ = [
    "MaskND",
]


def _resolve_ndim(
    ndim: int | None,
    coordinates: CoordinateSystemND | None,
    minimum_ndim: int | None = None,
) -> int | None:
    """Infer the target mask dimension from explicit or coordinate context."""

    if ndim is not None:
        return int(ndim)
    if coordinates is not None:
        return coordinates.ndim
    if minimum_ndim is not None:
        return int(minimum_ndim)
    return None


@dataclass(slots=True)
class _FunctionMaskND(MaskND):
    """Private mask implementation backed by ``fn(*coords)``."""

    fn: Callable[..., np.ndarray]

    def __init__(
        self,
        fn,
        ndim: int | None = None,
        coordinates: CoordinateSystemND | None = None,
    ):
        self.fn = fn
        MaskND.__init__(self, ndim=_resolve_ndim(ndim, coordinates), coordinates=coordinates)

    def _call_sample(self, sample: CoordinateSampleND) -> np.ndarray:
        coords = tuple(np.asarray(sample.coords[..., index]) for index in range(self.ndim))
        value = np.asarray(self.fn(*coords))
        try:
            return np.broadcast_to(value, sample.coords.shape[:-1])
        except ValueError as exc:
            raise ValueError(
                f"Callable mask result with shape {value.shape} cannot be broadcast to "
                f"{sample.coords.shape[:-1]}."
            ) from exc


@dataclass(slots=True)
class _ConstantMaskND(MaskND):
    """Private constant mask used for scalar coercion in mask algebra."""

    value: float | np.ndarray

    def __init__(
        self,
        value: float | np.ndarray,
        ndim: int | None = None,
        coordinates: CoordinateSystemND | None = None,
    ):
        self.value = value
        MaskND.__init__(self, ndim=_resolve_ndim(ndim, coordinates), coordinates=coordinates)

    def _call_sample(self, sample: CoordinateSampleND) -> np.ndarray:
        value = np.asarray(self.value, dtype=np.result_type(sample.coords, np.float64))
        try:
            return np.broadcast_to(value, sample.coords.shape[:-1])
        except ValueError as exc:
            raise ValueError(
                "Constant mask value with shape "
                f"{value.shape} cannot be broadcast to {sample.coords.shape[:-1]}."
            ) from exc


def ensure_masknd(value, ndim: int) -> MaskND:
    """Return ``value`` as a mask on ``R^ndim``.

    Existing :class:`MaskND` instances are validated for matching dimension.
    Scalars and broadcastable arrays are wrapped in an internal constant mask.
    """

    if isinstance(value, MaskND):
        if value.ndim != ndim:
            raise ValueError(f"Mask ndim {value.ndim} does not match required ndim {ndim}.")
        return value
    return _ConstantMaskND(value, ndim=ndim)
