from __future__ import annotations

import copy
import sys
from abc import abstractmethod
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Generic, cast, overload

import numpy as np
import pytest
from packaging.version import Version

from xarray.core.indexing import ExplicitlyIndexed
from xarray.namedarray._typing import (
    _arrayfunction_or_api,
    _default,
    _DType_co,
    _ShapeType_co,
)
from xarray.namedarray.core import NamedArray, from_array
from xarray.namedarray.utils import fake_target_chunksize
from xarray.tests import requires_cftime

if TYPE_CHECKING:
    from types import ModuleType

    from numpy.typing import ArrayLike, DTypeLike, NDArray

    from xarray.namedarray._typing import (
        Default,
        DuckArray,
        _AttrsLike,
        _Dim,
        _DimsLike,
        _DType,
        _IndexKeyLike,
        _IntOrUnknown,
        _Shape,
        _ShapeLike,
        duckarray,
    )


class CustomArrayBase(Generic[_ShapeType_co, _DType_co]):
    def __init__(self, array: duckarray[Any, _DType_co]) -> None:
        self.array: duckarray[Any, _DType_co] = array

    @property
    def dtype(self) -> _DType_co:
        return self.array.dtype

    @property
    def shape(self) -> _Shape:
        return self.array.shape


class CustomArray(
    CustomArrayBase[_ShapeType_co, _DType_co], Generic[_ShapeType_co, _DType_co]
):
    def __array__(
        self, dtype: DTypeLike | None = None, /, *, copy: bool | None = None
    ) -> np.ndarray[Any, np.dtype[np.generic]]:
        if Version(np.__version__) >= Version("2.0.0"):
            return np.asarray(self.array, dtype=dtype, copy=copy)
        else:
            return np.asarray(self.array, dtype=dtype)


class CustomArrayIndexable(
    CustomArrayBase[_ShapeType_co, _DType_co],
    ExplicitlyIndexed,
    Generic[_ShapeType_co, _DType_co],
):
    def __getitem__(
        self, key: _IndexKeyLike | CustomArrayIndexable[Any, Any], /
    ) -> CustomArrayIndexable[Any, _DType_co]:
        if isinstance(key, CustomArrayIndexable):
            if isinstance(key.array, type(self.array)):
                # TODO: key.array is duckarray here, can it be narrowed down further?
                # an _arrayapi cannot be used on a _arrayfunction for example.
                return type(self)(array=self.array[key.array])  # type: ignore[index]
            else:
                raise TypeError("key must have the same array type as self")
        else:
            return type(self)(array=self.array[key])

    def __array_namespace__(self) -> ModuleType:
        return np


def check_duck_array_typevar(a: duckarray[Any, _DType]) -> duckarray[Any, _DType]:
    # Mypy checks a is valid:
    b: duckarray[Any, _DType] = a

    # Runtime check if valid:
    if isinstance(b, _arrayfunction_or_api):
        return b
    else:
        missing_attrs = ""
        actual_attrs = set(dir(b))
        for t in _arrayfunction_or_api:
            if sys.version_info >= (3, 13):
                # https://github.com/python/cpython/issues/104873
                from typing import get_protocol_members

                expected_attrs = get_protocol_members(t)
            elif sys.version_info >= (3, 12):
                expected_attrs = t.__protocol_attrs__
            else:
                from typing import _get_protocol_attrs  # type: ignore[attr-defined]

                expected_attrs = _get_protocol_attrs(t)

            missing_attrs_ = expected_attrs - actual_attrs
            if missing_attrs_:
                missing_attrs += f"{t.__name__} - {missing_attrs_}\n"
        raise TypeError(
            f"a ({type(a)}) is not a valid _arrayfunction or _arrayapi. "
            "Missing following attrs:\n"
            f"{missing_attrs}"
        )


class NamedArraySubclassobjects:
    @pytest.fixture
    def target(self, data: np.ndarray[Any, Any]) -> Any:
        """Fixture that needs to be overridden"""
        raise NotImplementedError

    @abstractmethod
    def cls(self, *args: Any, **kwargs: Any) -> Any:
        """Method that needs to be overridden"""
        raise NotImplementedError

    @pytest.fixture
    def data(self) -> np.ndarray[Any, np.dtype[Any]]:
        return 0.5 * np.arange(10).reshape(2, 5)

    @pytest.fixture
    def random_inputs(self) -> np.ndarray[Any, np.dtype[np.float32]]:
        return np.arange(3 * 4 * 5, dtype=np.float32).reshape((3, 4, 5))

    def test_properties(self, target: Any, data: Any) -> None:
        assert target.dims == ("x", "y")
        assert np.array_equal(target.data, data)
        assert target.dtype == float
        assert target.shape == (2, 5)
        assert target.ndim == 2
        assert target.sizes == {"x": 2, "y": 5}
        assert target.size == 10
        assert target.nbytes == 80
        assert len(target) == 2

    def test_attrs(self, target: Any) -> None:
        assert target.attrs == {}
        attrs = {"foo": "bar"}
        target.attrs = attrs
        assert target.attrs == attrs
        assert isinstance(target.attrs, dict)
        target.attrs["foo"] = "baz"
        assert target.attrs["foo"] == "baz"

    @pytest.mark.parametrize(
        "expected", [np.array([1, 2], dtype=np.dtype(np.int8)), [1, 2]]
    )
    def test_init(self, expected: Any) -> None:
        actual = self.cls(("x",), expected)
        assert np.array_equal(np.asarray(actual.data), expected)

        actual = self.cls(("x",), expected)
        assert np.array_equal(np.asarray(actual.data), expected)

    def test_data(self, random_inputs: Any) -> None:
        expected = self.cls(["x", "y", "z"], random_inputs)
        assert np.array_equal(np.asarray(expected.data), random_inputs)
        with pytest.raises(ValueError):
            expected.data = np.random.random((3, 4)).astype(np.float64)
        d2 = np.arange(3 * 4 * 5, dtype=np.float32).reshape((3, 4, 5))
        expected.data = d2
        assert np.array_equal(np.asarray(expected.data), d2)


class TestNamedArray(NamedArraySubclassobjects):
    def cls(self, *args: Any, **kwargs: Any) -> NamedArray[Any, Any]:
        return NamedArray(*args, **kwargs)

    @pytest.fixture
    def target(self, data: np.ndarray[Any, Any]) -> NamedArray[Any, Any]:
        return NamedArray(["x", "y"], data)

    @pytest.mark.parametrize(
        "expected",
        [
            np.array([1, 2], dtype=np.dtype(np.int8)),
            pytest.param(
                [1, 2],
                marks=pytest.mark.xfail(
                    reason="NamedArray only supports array-like objects"
                ),
            ),
        ],
    )
    def test_init(self, expected: Any) -> None:
        super().test_init(expected)

    @pytest.mark.parametrize(
        "dims, data, expected, raise_error",
        [
            (("x",), [1, 2, 3], np.array([1, 2, 3]), False),
            ((1,), np.array([4, 5, 6]), np.array([4, 5, 6]), False),
            ((), 2, np.array(2), False),
            # Fail:
            (
                ("x",),
                NamedArray("time", np.array([1, 2, 3], dtype=np.dtype(np.int64))),
                np.array([1, 2, 3]),
                True,
            ),
        ],
    )
    def test_from_array(
        self,
        dims: _DimsLike,
        data: ArrayLike,
        expected: np.ndarray[Any, Any],
        raise_error: bool,
    ) -> None:
        actual: NamedArray[Any, Any]
        if raise_error:
            with pytest.raises(TypeError, match="already a Named array"):
                actual = from_array(dims, data)

                # Named arrays are not allowed:
                from_array(actual)  # type: ignore[call-overload]
        else:
            actual = from_array(dims, data)

            assert np.array_equal(np.asarray(actual.data), expected)

    def test_from_array_with_masked_array(self) -> None:
        masked_array: np.ndarray[Any, np.dtype[np.generic]]
        masked_array = np.ma.array([1, 2, 3], mask=[False, True, False])  # type: ignore[no-untyped-call]
        with pytest.raises(NotImplementedError):
            from_array(("x",), masked_array)

    def test_from_array_with_0d_object(self) -> None:
        data = np.empty((), dtype=object)
        data[()] = (10, 12, 12)
        narr = from_array((), data)
        np.array_equal(np.asarray(narr.data), data)

    # TODO: Make xr.core.indexing.ExplicitlyIndexed pass as a subclass of_arrayfunction_or_api
    # and remove this test.
    def test_from_array_with_explicitly_indexed(
        self, random_inputs: np.ndarray[Any, Any]
    ) -> None:
        array: CustomArray[Any, Any]
        array = CustomArray(random_inputs)
        output: NamedArray[Any, Any]
        output = from_array(("x", "y", "z"), array)
        assert isinstance(output.data, np.ndarray)

        array2: CustomArrayIndexable[Any, Any]
        array2 = CustomArrayIndexable(random_inputs)
        output2: NamedArray[Any, Any]
        output2 = from_array(("x", "y", "z"), array2)
        assert isinstance(output2.data, CustomArrayIndexable)

    def test_real_and_imag(self) -> None:
        expected_real: np.ndarray[Any, np.dtype[np.float64]]
        expected_real = np.arange(3, dtype=np.float64)

        expected_imag: np.ndarray[Any, np.dtype[np.float64]]
        expected_imag = -np.arange(3, dtype=np.float64)

        arr: np.ndarray[Any, np.dtype[np.complex128]]
        arr = expected_real + 1j * expected_imag

        named_array: NamedArray[Any, np.dtype[np.complex128]]
        named_array = NamedArray(["x"], arr)

        actual_real: duckarray[Any, np.dtype[np.float64]] = named_array.real.data
        assert np.array_equal(np.asarray(actual_real), expected_real)
        assert actual_real.dtype == expected_real.dtype

        actual_imag: duckarray[Any, np.dtype[np.float64]] = named_array.imag.data
        assert np.array_equal(np.asarray(actual_imag), expected_imag)
        assert actual_imag.dtype == expected_imag.dtype

    # Additional tests as per your original class-based code
    @pytest.mark.parametrize(
        "data, dtype",
        [
            ("foo", np.dtype("U3")),
            (b"foo", np.dtype("S3")),
        ],
    )
    def test_from_array_0d_string(self, data: Any, dtype: DTypeLike | None) -> None:
        named_array: NamedArray[Any, Any]
        named_array = from_array([], data)
        assert named_array.data == data
        assert named_array.dims == ()
        assert named_array.sizes == {}
        assert named_array.attrs == {}
        assert named_array.ndim == 0
        assert named_array.size == 1
        assert named_array.dtype == dtype

    def test_from_array_0d_object(self) -> None:
        named_array: NamedArray[Any, Any]
        named_array = from_array([], (10, 12, 12))
        expected_data = np.empty((), dtype=object)
        expected_data[()] = (10, 12, 12)
        assert np.array_equal(np.asarray(named_array.data), expected_data)

        assert named_array.dims == ()
        assert named_array.sizes == {}
        assert named_array.attrs == {}
        assert named_array.ndim == 0
        assert named_array.size == 1
        assert named_array.dtype == np.dtype("O")

    def test_from_array_0d_datetime(self) -> None:
        named_array: NamedArray[Any, Any]
        named_array = from_array([], np.datetime64("2000-01-01"))
        assert named_array.dtype == np.dtype("datetime64[D]")

    @pytest.mark.parametrize(
        "timedelta, expected_dtype",
        [
            (np.timedelta64(1, "D"), np.dtype("timedelta64[D]")),
            (np.timedelta64(1, "s"), np.dtype("timedelta64[s]")),
            (np.timedelta64(1, "m"), np.dtype("timedelta64[m]")),
            (np.timedelta64(1, "h"), np.dtype("timedelta64[h]")),
            (np.timedelta64(1, "us"), np.dtype("timedelta64[us]")),
            (np.timedelta64(1, "ns"), np.dtype("timedelta64[ns]")),
            (np.timedelta64(1, "ps"), np.dtype("timedelta64[ps]")),
            (np.timedelta64(1, "fs"), np.dtype("timedelta64[fs]")),
            (np.timedelta64(1, "as"), np.dtype("timedelta64[as]")),
        ],
    )
    def test_from_array_0d_timedelta(
        self, timedelta: np.timedelta64, expected_dtype: np.dtype[np.timedelta64]
    ) -> None:
        named_array: NamedArray[Any, Any]
        named_array = from_array([], timedelta)
        assert named_array.dtype == expected_dtype
        assert named_array.data == timedelta

    @pytest.mark.parametrize(
        "dims, data_shape, new_dims, raises",
        [
            (["x", "y", "z"], (2, 3, 4), ["a", "b", "c"], False),
            (["x", "y", "z"], (2, 3, 4), ["a", "b"], True),
            (["x", "y", "z"], (2, 4, 5), ["a", "b", "c", "d"], True),
            ([], [], (), False),
            ([], [], ("x",), True),
        ],
    )
    def test_dims_setter(
        self, dims: Any, data_shape: Any, new_dims: Any, raises: bool
    ) -> None:
        named_array: NamedArray[Any, Any]
        named_array = NamedArray(dims, np.asarray(np.random.random(data_shape)))
        assert named_array.dims == tuple(dims)
        if raises:
            with pytest.raises(ValueError):
                named_array.dims = new_dims
        else:
            named_array.dims = new_dims
            assert named_array.dims == tuple(new_dims)

    def test_duck_array_class(self) -> None:
        numpy_a: NDArray[np.int64]
        numpy_a = np.array([2.1, 4], dtype=np.dtype(np.int64))
        check_duck_array_typevar(numpy_a)

        masked_a: np.ma.MaskedArray[Any, np.dtype[np.int64]]
        masked_a = np.ma.asarray([2.1, 4], dtype=np.dtype(np.int64))  # type: ignore[no-untyped-call]
        check_duck_array_typevar(masked_a)  # type: ignore[arg-type]  # MaskedArray not in duckarray union

        custom_a: CustomArrayIndexable[Any, np.dtype[np.int64]]
        custom_a = CustomArrayIndexable(numpy_a)
        check_duck_array_typevar(custom_a)

    def test_duck_array_class_array_api(self) -> None:
        # Test numpy's array api:
        nxp = pytest.importorskip("array_api_strict", minversion="1.0")

        # TODO: nxp doesn't use dtype typevars, so can only use Any for the moment:
        arrayapi_a: duckarray[Any, Any]  #  duckarray[Any, np.dtype[np.int64]]
        arrayapi_a = nxp.asarray([2.1, 4], dtype=nxp.int64)
        check_duck_array_typevar(arrayapi_a)

    def test_new_namedarray(self) -> None:
        dtype_float = np.dtype(np.float32)
        narr_float: NamedArray[Any, np.dtype[np.float32]]
        narr_float = NamedArray(("x",), np.array([1.5, 3.2], dtype=dtype_float))
        assert narr_float.dtype == dtype_float

        dtype_int = np.dtype(np.int8)
        narr_int: NamedArray[Any, np.dtype[np.int8]]
        narr_int = narr_float._new(("x",), np.array([1, 3], dtype=dtype_int))
        assert narr_int.dtype == dtype_int

        class Variable(
            NamedArray[_ShapeType_co, _DType_co], Generic[_ShapeType_co, _DType_co]
        ):
            @overload
            def _new(
                self,
                dims: _DimsLike | Default = ...,
                data: duckarray[Any, _DType] = ...,
                attrs: _AttrsLike | Default = ...,
            ) -> Variable[Any, _DType]: ...

            @overload
            def _new(
                self,
                dims: _DimsLike | Default = ...,
                data: Default = ...,
                attrs: _AttrsLike | Default = ...,
            ) -> Variable[_ShapeType_co, _DType_co]: ...

            def _new(
                self,
                dims: _DimsLike | Default = _default,
                data: duckarray[Any, _DType] | Default = _default,
                attrs: _AttrsLike | Default = _default,
            ) -> Variable[Any, _DType] | Variable[_ShapeType_co, _DType_co]:
                dims_ = copy.copy(self._dims) if dims is _default else dims

                attrs_: Mapping[Any, Any] | None
                if attrs is _default:
                    attrs_ = None if self._attrs is None else self._attrs.copy()
                else:
                    attrs_ = attrs

                if data is _default:
                    return type(self)(dims_, copy.copy(self._data), attrs_)
                cls_ = cast("type[Variable[Any, _DType]]", type(self))
                return cls_(dims_, data, attrs_)

        var_float: Variable[Any, np.dtype[np.float32]]
        var_float = Variable(("x",), np.array([1.5, 3.2], dtype=dtype_float))
        assert var_float.dtype == dtype_float

        var_int: Variable[Any, np.dtype[np.int8]]
        var_int = var_float._new(("x",), np.array([1, 3], dtype=dtype_int))
        assert var_int.dtype == dtype_int

    def test_replace_namedarray(self) -> None:
        dtype_float = np.dtype(np.float32)
        np_val: np.ndarray[Any, np.dtype[np.float32]]
        np_val = np.array([1.5, 3.2], dtype=dtype_float)
        np_val2: np.ndarray[Any, np.dtype[np.float32]]
        np_val2 = 2 * np_val

        narr_float: NamedArray[Any, np.dtype[np.float32]]
        narr_float = NamedArray(("x",), np_val)
        assert narr_float.dtype == dtype_float

        narr_float2: NamedArray[Any, np.dtype[np.float32]]
        narr_float2 = NamedArray(("x",), np_val2)
        assert narr_float2.dtype == dtype_float

        class Variable(
            NamedArray[_ShapeType_co, _DType_co], Generic[_ShapeType_co, _DType_co]
        ):
            @overload
            def _new(
                self,
                dims: _DimsLike | Default = ...,
                data: duckarray[Any, _DType] = ...,
                attrs: _AttrsLike | Default = ...,
            ) -> Variable[Any, _DType]: ...

            @overload
            def _new(
                self,
                dims: _DimsLike | Default = ...,
                data: Default = ...,
                attrs: _AttrsLike | Default = ...,
            ) -> Variable[_ShapeType_co, _DType_co]: ...

            def _new(
                self,
                dims: _DimsLike | Default = _default,
                data: duckarray[Any, _DType] | Default = _default,
                attrs: _AttrsLike | Default = _default,
            ) -> Variable[Any, _DType] | Variable[_ShapeType_co, _DType_co]:
                dims_ = copy.copy(self._dims) if dims is _default else dims

                attrs_: Mapping[Any, Any] | None
                if attrs is _default:
                    attrs_ = None if self._attrs is None else self._attrs.copy()
                else:
                    attrs_ = attrs

                if data is _default:
                    return type(self)(dims_, copy.copy(self._data), attrs_)
                cls_ = cast("type[Variable[Any, _DType]]", type(self))
                return cls_(dims_, data, attrs_)

        var_float: Variable[Any, np.dtype[np.float32]]
        var_float = Variable(("x",), np_val)
        assert var_float.dtype == dtype_float

        var_float2: Variable[Any, np.dtype[np.float32]]
        var_float2 = var_float._replace(("x",), np_val2)
        assert var_float2.dtype == dtype_float

    @pytest.mark.parametrize(
        "dim,expected_ndim,expected_shape,expected_dims",
        [
            (None, 3, (1, 2, 5), (None, "x", "y")),
            (_default, 3, (1, 2, 5), ("dim_2", "x", "y")),
            ("z", 3, (1, 2, 5), ("z", "x", "y")),
        ],
    )
    def test_expand_dims(
        self,
        target: NamedArray[Any, np.dtype[np.float32]],
        dim: _Dim | Default,
        expected_ndim: int,
        expected_shape: _ShapeLike,
        expected_dims: _DimsLike,
    ) -> None:
        result = target.expand_dims(dim=dim)
        assert result.ndim == expected_ndim
        assert result.shape == expected_shape
        assert result.dims == expected_dims

    @pytest.mark.parametrize(
        "dims, expected_sizes",
        [
            ((), {"y": 5, "x": 2}),
            (["y", "x"], {"y": 5, "x": 2}),
            (["y", ...], {"y": 5, "x": 2}),
        ],
    )
    def test_permute_dims(
        self,
        target: NamedArray[Any, np.dtype[np.float32]],
        dims: _DimsLike,
        expected_sizes: dict[_Dim, _IntOrUnknown],
    ) -> None:
        actual = target.permute_dims(*dims)
        assert actual.sizes == expected_sizes

    def test_permute_dims_errors(
        self,
        target: NamedArray[Any, np.dtype[np.float32]],
    ) -> None:
        with pytest.raises(ValueError, match=r"'y'.*permuted list"):
            dims = ["y"]
            target.permute_dims(*dims)

    @pytest.mark.parametrize(
        "broadcast_dims,expected_ndim",
        [
            ({"x": 2, "y": 5}, 2),
            ({"x": 2, "y": 5, "z": 2}, 3),
            ({"w": 1, "x": 2, "y": 5}, 3),
        ],
    )
    def test_broadcast_to(
        self,
        target: NamedArray[Any, np.dtype[np.float32]],
        broadcast_dims: Mapping[_Dim, int],
        expected_ndim: int,
    ) -> None:
        expand_dims = set(broadcast_dims.keys()) - set(target.dims)
        # loop over expand_dims and call .expand_dims(dim=dim) in a loop
        for dim in expand_dims:
            target = target.expand_dims(dim=dim)
        result = target.broadcast_to(broadcast_dims)
        assert result.ndim == expected_ndim
        assert result.sizes == broadcast_dims

    def test_broadcast_to_errors(
        self, target: NamedArray[Any, np.dtype[np.float32]]
    ) -> None:
        with pytest.raises(
            ValueError,
            match=r"operands could not be broadcast together with remapped shapes",
        ):
            target.broadcast_to({"x": 2, "y": 2})

        with pytest.raises(ValueError, match=r"Cannot add new dimensions"):
            target.broadcast_to({"x": 2, "y": 2, "z": 2})

    def test_warn_on_repeated_dimension_names(self) -> None:
        with pytest.warns(UserWarning, match="Duplicate dimension names"):
            NamedArray(("x", "x"), np.arange(4).reshape(2, 2))

    def test_aggregation(self) -> None:
        x: NamedArray[Any, np.dtype[np.int64]]
        x = NamedArray(("x", "y"), np.arange(4).reshape(2, 2))

        result = x.sum()
        assert isinstance(result.data, np.ndarray)


def test_repr() -> None:
    x: NamedArray[Any, np.dtype[np.uint64]]
    x = NamedArray(("x",), np.array([0], dtype=np.uint64))

    # Reprs should not crash:
    r = x.__repr__()
    x._repr_html_()

    # Basic comparison:
    assert r == "<xarray.NamedArray (x: 1)> Size: 8B\narray([0], dtype=uint64)"


@pytest.mark.parametrize(
    "input_array, expected_chunksize_faked, expected_dtype",
    [
        (np.arange(100).reshape(10, 10), 1024, np.int64),
        (np.arange(100).reshape(10, 10).astype(np.float32), 1024, np.float32),
    ],
)
def test_fake_target_chunksize(
    input_array: DuckArray[Any],
    expected_chunksize_faked: int,
    expected_dtype: DTypeLike,
) -> None:
    """
    Check that `fake_target_chunksize` returns the expected chunksize and dtype.
    - It pretends to dask we are chunking an array with an 8-byte dtype, ie. a float64.
    As such, it will *double* the amount of memory a 4-byte dtype (like float32) would try to use,
    fooling it into actually using the correct amount of memory. For object dtypes, which are
    generally larger, it will reduce the effective dask configuration chunksize, reducing the size of
    the arrays per chunk such that we get the same amount of memory used.
    """
    target_chunksize = 1024

    faked_chunksize, dtype = fake_target_chunksize(input_array, target_chunksize)

    assert faked_chunksize == expected_chunksize_faked
    assert dtype == expected_dtype


@requires_cftime
def test_fake_target_chunksize_cftime() -> None:
    """
    Check that `fake_target_chunksize` returns the expected chunksize and dtype.
    - It pretends to dask we are chunking an array with an 8-byte dtype, ie. a float64.
    - This is the same as the above test, but specifically for a CFTime array case - split for testing reasons
    """
    import cftime

    target_chunksize = 1024

    input_array = np.array(
        [
            cftime.Datetime360Day(2000, month, day, 0, 0, 0, 0)
            for month in range(1, 11)
            for day in range(1, 11)
        ],
        dtype=object,
    ).reshape(10, 10)

    faked_chunksize, dtype = fake_target_chunksize(input_array, target_chunksize)  # type: ignore[arg-type,unused-ignore]

    assert faked_chunksize == 73
    assert dtype == np.float64