# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE

from __future__ import annotations

from awkward._nplikes.array_like import MaterializableArray, maybe_materialize
from awkward._nplikes.numpy_like import NumpyLike, NumpyMetadata
from awkward._nplikes.shape import ShapeItem, unknown_length
from awkward._typing import TYPE_CHECKING, Any, DType, Self

np = NumpyMetadata.instance()

if TYPE_CHECKING:
    from numpy.typing import DTypeLike


class PlaceholderArray(MaterializableArray):
    def __init__(
        self,
        nplike: NumpyLike,
        shape: tuple[ShapeItem, ...],
        dtype: DType,
        buffer_key: str | None = None,
    ):
        self._nplike = nplike
        self._shape = tuple(dim if dim is unknown_length else int(dim) for dim in shape)
        self._dtype = np.dtype(dtype)
        self._buffer_key = buffer_key

    @property
    def buffer_key(self) -> str | None:
        return self._buffer_key

    @property
    def dtype(self) -> DType:
        return self._dtype

    @property
    def shape(self) -> tuple[ShapeItem, ...]:
        return self._shape

    @property
    def ndim(self) -> int:
        return len(self._shape)

    @property
    def size(self) -> ShapeItem:
        size: ShapeItem = 1
        for item in self._shape:
            size *= item
        return size

    @property
    def nbytes(self) -> int:
        return 0

    def materialize(self):
        msg = f"{self} should never have been encountered."
        if self.buffer_key:
            msg += (
                f" Awkward Array tried to access a buffer at '{self.buffer_key}', "
                "but it exists only as a placeholder."
            )
        msg += (
            " This is unexpected behavior — please open an issue at "
            "https://github.com/scikit-hep/awkward/issues with a minimal example."
        )
        raise RuntimeError(msg)

    @property
    def strides(self) -> tuple[ShapeItem, ...]:
        out: tuple[ShapeItem, ...] = (self._dtype.itemsize,)
        for item in reversed(self._shape):
            out = (item * out[0], *out)
        return out

    @property
    def T(self):
        return type(self)(self._nplike, self._shape[::-1], self._dtype)

    def view(self, dtype: DTypeLike) -> Self:
        dtype = np.dtype(dtype)
        if len(self._shape) >= 1:
            last, remainder = divmod(
                self._shape[-1] * self._dtype.itemsize, dtype.itemsize
            )
            if remainder is not unknown_length and remainder != 0:
                raise ValueError(
                    "new size of array with larger dtype must be a "
                    "divisor of the total size in bytes (of the last axis of the array)"
                )
            shape = (*self._shape[:-1], last)
        else:
            shape = self._shape
        return type(self)(self._nplike, shape, dtype, self._buffer_key)

    def __repr__(self):
        dtype = repr(self._dtype)
        shape = f", shape={self._shape!r}"
        return f"PlaceholderArray({dtype}{shape})"

    def __getitem__(self, index):
        # Typetracers permit slices that don't touch data or shapes
        if isinstance(index, slice):
            length = self._shape[0]

            # Unknown-length placeholders should not be sliced (as their shapes would be touched(
            if length is unknown_length:
                raise AssertionError(
                    "placeholder arrays that are sliced should have known shapes"
                )
            # Known-length placeholders *always* need a known shape
            elif (
                index.start is unknown_length
                or index.stop is unknown_length
                or index.step is unknown_length
            ):
                raise AssertionError(
                    "known-length placeholders should never encounter unknown lengths in slices"
                )
            else:
                start, stop, step = index.indices(length)
                new_length = max(
                    0, (stop - start + (step - (1 if step > 0 else -1))) // step
                )

            return type(self)(
                self._nplike, (new_length,), self._dtype, self._buffer_key
            )
        else:
            msg = f"{type(self).__name__} supports only trivial slices, not {type(index).__name__}"
            if self.buffer_key:
                msg += f"\n\nAwkward-array attempted to access a buffer at '{self.buffer_key}', but "
                msg += (
                    "it has been excluded during a pre-run phase (possibly by Dask). "
                )
                msg += "If this was supposed to happen automatically (e.g. you're using Dask), "
                msg += "please report it to the developers at: https://github.com/scikit-hep/awkward/issues"
            raise TypeError(msg)

    def tolist(self):
        self.materialize()

    def __setitem__(self, key, value):
        del key
        maybe_materialize(self, value)

    def __bool__(self):  # pylint: disable=E0304
        self.materialize()

    def __int__(self):
        self.materialize()

    def __index__(self):  # noqa: PLE0305 # pylint: disable=E0305
        self.materialize()

    def __len__(self) -> int:
        return int(self._shape[0])

    def __add__(self, other):
        maybe_materialize(self, other)

    def __and__(self, other):
        maybe_materialize(self, other)

    def __eq__(self, other):
        maybe_materialize(self, other)

    def __floordiv__(self, other):
        maybe_materialize(self, other)

    def __ge__(self, other):
        maybe_materialize(self, other)

    def __gt__(self, other):
        maybe_materialize(self, other)

    def __invert__(self):
        self.materialize()

    def __le__(self, other):
        maybe_materialize(self, other)

    def __lt__(self, other):
        maybe_materialize(self, other)

    def __mul__(self, other):
        maybe_materialize(self, other)

    def __or__(self, other):
        maybe_materialize(self, other)

    def __sub__(self, other):
        maybe_materialize(self, other)

    def __truediv__(self, other):
        maybe_materialize(self, other)

    def __iter__(self):  # pylint: disable=E0301
        self.materialize()

    def __array__(self, dtype=None, copy=None):
        self.materialize()

    def __cupy_get_ndarray__(self):
        self.materialize()

    def __jax_array__(self):
        self.materialize()

    @property
    def __array_interface__(self):
        self.materialize()

    @property
    def __cuda_array_interface__(self):
        self.materialize()

    def __dlpack_device__(self):
        self.materialize()

    def __dlpack__(self, stream: Any = None):
        del stream
        self.materialize()
