"""Helpers for interfacing array like objects."""

import copy
import ctypes
import json
import warnings
from abc import ABC, abstractmethod
from functools import cache as fcache
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    List,
    Literal,
    Optional,
    Protocol,
    Tuple,
    Type,
    TypeAlias,
    TypedDict,
    TypeGuard,
    Union,
    cast,
    overload,
)

import numpy as np

from ._typing import (
    ArrowCatList,
    CNumericPtr,
    DataType,
    FeatureTypes,
    NumpyDType,
    NumpyOrCupy,
)
from .compat import import_cupy, import_pyarrow, lazy_isinstance

if TYPE_CHECKING:
    import pandas as pd
    import pyarrow as pa


# Used for accepting inputs for numpy and cupy arrays
class _ArrayLikeArg(Protocol):
    @property
    def __array_interface__(self) -> "ArrayInf": ...


class _CudaArrayLikeArg(Protocol):
    @property
    def __cuda_array_interface__(self) -> "CudaArrayInf": ...


ArrayInf = TypedDict(
    "ArrayInf",
    {
        "data": Tuple[int, bool],
        "typestr": str,
        "version": Literal[3],
        "strides": Optional[Tuple[int, ...]],
        "shape": Tuple[int, ...],
        "mask": Union["ArrayInf", None, _ArrayLikeArg],
    },
)

CudaArrayInf = TypedDict(
    "CudaArrayInf",
    {
        "data": Tuple[int, bool],
        "typestr": str,
        "version": Literal[3],
        "strides": Optional[Tuple[int, ...]],
        "shape": Tuple[int, ...],
        "mask": Union["ArrayInf", None, _ArrayLikeArg],
        "stream": int,
    },
)

StringArray = TypedDict("StringArray", {"offsets": ArrayInf, "values": ArrayInf})
CudaStringArray = TypedDict(
    "CudaStringArray", {"offsets": CudaArrayInf, "values": CudaArrayInf}
)


def array_hasobject(data: DataType) -> bool:
    """Whether the numpy array has object dtype."""
    return (
        hasattr(data, "dtype")
        and hasattr(data.dtype, "hasobject")
        and data.dtype.hasobject
    )


def cuda_array_interface_dict(data: _CudaArrayLikeArg) -> CudaArrayInf:
    """Returns a dictionary storing the CUDA array interface."""
    if array_hasobject(data):
        raise ValueError("Input data contains `object` dtype.  Expecting numeric data.")
    ainf = data.__cuda_array_interface__
    if "mask" in ainf and ainf["mask"] is not None:
        mask_ainf = ainf["mask"].__cuda_array_interface__  # type: ignore[union-attr]
        # Normalize the validity mask to XGBoost's expected layout (`|t1` bit field of
        # length `n_samples`).
        typestr = mask_ainf["typestr"]
        n_samples = ainf["shape"][0]
        if typestr[1] in ("u", "i", "t") and typestr[2:] == "1" and n_samples:
            mask_ainf = dict(mask_ainf)
            mask_ainf["typestr"] = "|t1"
            mask_ainf["shape"] = (n_samples,)
            mask_ainf.pop("strides", None)
        ainf["mask"] = mask_ainf
    return ainf


def cuda_array_interface(data: _CudaArrayLikeArg) -> bytes:
    """Make cuda array interface str."""
    interface = cuda_array_interface_dict(data)
    interface_str = bytes(json.dumps(interface), "utf-8")
    return interface_str


def from_array_interface(interface: ArrayInf, zero_copy: bool = False) -> NumpyOrCupy:
    """Convert array interface to numpy or cupy array"""

    class Array:
        """Wrapper type for communicating with numpy and cupy."""

        _interface: Optional[ArrayInf] = None

        @property
        def __array_interface__(self) -> Optional[ArrayInf]:
            return self._interface

        @__array_interface__.setter
        def __array_interface__(self, interface: ArrayInf) -> None:
            self._interface = copy.copy(interface)
            # Convert some fields to tuple as required by numpy
            self._interface["shape"] = tuple(self._interface["shape"])
            self._interface["data"] = (
                self._interface["data"][0],
                self._interface["data"][1],
            )
            strides = self._interface.get("strides", None)
            if strides is not None:
                self._interface["strides"] = tuple(strides)

        @property
        def __cuda_array_interface__(self) -> Optional[ArrayInf]:
            return self.__array_interface__

        @__cuda_array_interface__.setter
        def __cuda_array_interface__(self, interface: ArrayInf) -> None:
            self.__array_interface__ = interface

        @property
        def shape(self) -> Tuple[int, ...]:
            """Shape of the input array."""
            aif = self.__array_interface__
            assert aif is not None
            return aif["shape"]

        @property
        def size(self) -> np.signedinteger:
            """Total size of the input array."""
            return np.prod(self.shape)

    arr = Array()

    # Cupy and numpy might run into issue when constructing an empty array from an array
    # interface. we explicitly check for emptiness.
    if "stream" in interface:
        # CUDA stream is presented, this is a __cuda_array_interface__.
        arr.__cuda_array_interface__ = interface
        cp = import_cupy()
        if arr.size == 0:
            return cp.empty(shape=arr.shape, dtype=np.dtype(interface["typestr"]))
        out = cp.array(arr, copy=not zero_copy)
    else:
        arr.__array_interface__ = interface
        if arr.size == 0:
            return np.empty(shape=arr.shape, dtype=np.dtype(interface["typestr"]))
        out = np.array(arr, copy=not zero_copy)

    return out


# Default constant value for CUDA per-thread stream.
STREAM_PER_THREAD = 2


# Typing is not strict as there are subtle differences between CUDA array interface and
# array interface. We handle them uniformly for now.
def make_array_interface(
    ptr: Union[CNumericPtr, int],
    shape: Tuple[int, ...],
    dtype: Type[np.number],
    is_cuda: bool,
) -> ArrayInf:
    """Make an __(cuda)_array_interface__ from a pointer."""
    # Use an empty array to handle typestr and descr
    if is_cuda:
        empty = import_cupy().empty(shape=(0,), dtype=dtype)
        array = empty.__cuda_array_interface__  # pylint: disable=no-member
    else:
        empty = np.empty(shape=(0,), dtype=dtype)
        array = empty.__array_interface__  # pylint: disable=no-member

    if not isinstance(ptr, int):
        addr = ctypes.cast(ptr, ctypes.c_void_p).value
    else:
        addr = ptr
    length = int(np.prod(shape))
    # Handle empty dataset.
    assert addr is not None or length == 0

    if addr is None:
        return array

    array["data"] = (addr, True)
    if is_cuda and "stream" not in array:
        array["stream"] = STREAM_PER_THREAD
    array["shape"] = shape
    array["strides"] = None
    return array


def is_arrow_dict(data: Any) -> TypeGuard["pa.DictionaryArray"]:
    """Is this an arrow dictionary array?"""
    return lazy_isinstance(data, "pyarrow.lib", "DictionaryArray")


class DfCatAccessor(Protocol):
    """Protocol for pandas cat accessor."""

    @property
    def categories(  # pylint: disable=missing-function-docstring
        self,
    ) -> "pd.Index": ...

    @property
    def codes(self) -> "pd.Series": ...  # pylint: disable=missing-function-docstring

    @property
    def dtype(self) -> np.dtype: ...  # pylint: disable=missing-function-docstring

    @property
    def values(self) -> np.ndarray: ...  # pylint: disable=missing-function-docstring

    def to_arrow(  # pylint: disable=missing-function-docstring
        self,
    ) -> Union["pa.StringArray", "pa.IntegerArray"]: ...

    @property
    def __cuda_array_interface__(self) -> CudaArrayInf: ...

    @property
    def _column(self) -> Any: ...


def _is_df_cat(data: Any) -> TypeGuard[DfCatAccessor]:
    # Test pd.Series.cat, not pd.Series
    return hasattr(data, "categories") and hasattr(data, "codes")


@fcache
def _arrow_npdtype() -> Dict[Any, Type[np.number]]:
    import pyarrow as pa

    mapping: Dict[Any, Type[np.number]] = {
        pa.int8(): np.int8,
        pa.int16(): np.int16,
        pa.int32(): np.int32,
        pa.int64(): np.int64,
        pa.uint8(): np.uint8,
        pa.uint16(): np.uint16,
        pa.uint32(): np.uint32,
        pa.uint64(): np.uint64,
        pa.float16(): np.float16,
        pa.float32(): np.float32,
        pa.float64(): np.float64,
    }

    return mapping


@overload
def _arrow_buf_inf(address: int, typestr: str, size: int, stream: None) -> ArrayInf: ...


@overload
def _arrow_buf_inf(
    address: int, typestr: str, size: int, stream: int
) -> CudaArrayInf: ...


def _arrow_buf_inf(
    address: int, typestr: str, size: int, stream: Optional[int]
) -> Union[ArrayInf, CudaArrayInf]:
    if stream is not None:
        jcuaif: CudaArrayInf = {
            "data": (address, True),
            "typestr": typestr,
            "version": 3,
            "strides": None,
            "shape": (size,),
            "mask": None,
            "stream": stream,
        }
        return jcuaif

    jaif: ArrayInf = {
        "data": (address, True),
        "typestr": typestr,
        "version": 3,
        "strides": None,
        "shape": (size,),
        "mask": None,
    }
    return jaif


def _arrow_cat_names_inf(cats: "pa.StringArray") -> Tuple[StringArray, Any]:
    if not TYPE_CHECKING:
        pa = import_pyarrow()

    # FIXME(jiamingy): Account for offset, need to find an implementation that returns
    # offset > 0
    assert cats.offset == 0
    buffers: List[pa.Buffer] = cats.buffers()
    mask, offset, data = buffers
    assert offset.is_cpu

    off_len = len(cats) + 1

    def get_n_bytes(typ: Type) -> int:
        return off_len * (np.iinfo(typ).bits // 8)

    if offset.size == get_n_bytes(np.int64):
        if not isinstance(cats, pa.LargeStringArray):
            arrow_str_error = "Expecting a `pyarrow.Array`."
            raise TypeError(arrow_str_error + f" Got: {type(cats)}.")
        # Convert to 32bit integer, arrow recommends against the use of i64. Also,
        # XGBoost cannot handle large number of categories (> 2**31).
        i32cats = cats.cast(pa.string())
        mask, offset, data = i32cats.buffers()

    if offset.size != get_n_bytes(np.int32):
        raise TypeError(
            "Arrow dictionary type offsets is required to be 32-bit integer."
        )

    joffset = _arrow_buf_inf(offset.address, "<i4", off_len, None)
    jdata = _arrow_buf_inf(data.address, "|i1", data.size, None)
    # Categories should not have missing values.
    assert mask is None

    jnames: StringArray = {"offsets": joffset, "values": jdata}
    return jnames, (mask, offset, data)


def _arrow_array_inf(
    array: "pa.Array",
) -> ArrayInf:
    """Helper for handling categorical codes."""
    if not TYPE_CHECKING:
        pa = import_pyarrow()
    if not isinstance(array, pa.Array):  # pylint: disable=E0606
        raise TypeError(f"Invalid input type: {type(array)}")

    mask, data = array.buffers()
    jdata = make_array_interface(
        data.address,
        shape=(len(array),),
        dtype=_arrow_npdtype()[array.type],
        is_cuda=not data.is_cpu,
    )

    if mask is not None:
        jmask: Optional[ArrayInf] = {
            "data": (mask.address, True),
            "typestr": "<t1",
            "version": 3,
            "strides": None,
            "shape": (len(array),),
            "mask": None,
        }
        if not mask.is_cpu:
            jmask["stream"] = STREAM_PER_THREAD  # type: ignore[index, typeddict-unknown-key]
    else:
        jmask = None

    jdata["mask"] = jmask
    return jdata


def arrow_cat_inf(  # pylint: disable=too-many-locals
    cats: "pa.StringArray",
    codes: Union[_ArrayLikeArg, _CudaArrayLikeArg, "pa.IntegerArray"],
) -> Tuple[StringArray, ArrayInf, Tuple]:
    """Get the array interface representation of a string-based category array."""
    jnames, cats_tmp = _arrow_cat_names_inf(cats)
    jcodes = _arrow_array_inf(codes)

    return jnames, jcodes, (cats_tmp, None)


def _ensure_np_dtype(
    data: DataType, dtype: Optional[NumpyDType]
) -> Tuple[np.ndarray, Optional[NumpyDType]]:
    """Ensure the np array has correct type and is contiguous."""
    if array_hasobject(data) or data.dtype in [np.float16, np.bool_]:
        dtype = np.float32
        data = data.astype(dtype, copy=False)
    if not data.flags.aligned:
        data = np.require(data, requirements="A")
    return data, dtype


def _is_flatten(array: NumpyOrCupy) -> bool:
    return len(array.shape) == 1 or array.shape[1] == 1


def array_interface_dict(data: np.ndarray) -> ArrayInf:
    """Returns an array interface from the input."""
    if array_hasobject(data):
        raise ValueError("Input data contains `object` dtype.  Expecting numeric data.")
    ainf = data.__array_interface__
    if "mask" in ainf:
        ainf["mask"] = ainf["mask"].__array_interface__
    return cast(ArrayInf, ainf)


def pd_cat_inf(  # pylint: disable=too-many-locals
    cats: DfCatAccessor, codes: "pd.Series"
) -> Tuple[Union[StringArray, ArrayInf], ArrayInf, Tuple]:
    """Get the array interface representation of pandas category accessor."""
    # pandas uses -1 to represent missing values for categorical features
    codes = codes.replace(-1, np.nan)

    def is_prim() -> bool:
        dtype = cats.dtype
        try:
            return np.issubdtype(dtype, np.floating) or np.issubdtype(dtype, np.integer)
        except TypeError:
            return False

    if is_prim():
        # Numeric index type
        name_values_num = cats.values
        jarr_values = array_interface_dict(name_values_num)
        code_values = codes.values
        jarr_codes = array_interface_dict(code_values)
        return jarr_values, jarr_codes, (name_values_num, code_values)

    def npstr_to_arrow_strarr(strarr: Any) -> Tuple[np.ndarray, str]:
        """Convert a string-like array to an arrow string array."""
        if not isinstance(strarr, np.ndarray):
            if hasattr(strarr, "to_numpy"):
                strarr = strarr.to_numpy(dtype=object)
            else:
                strarr = np.asarray(strarr, dtype=object)

        lenarr = np.vectorize(len)
        offsets = np.cumsum(
            np.concatenate([np.array([0], dtype=np.int64), lenarr(strarr)])
        )
        if strarr.dtype.kind == "S":
            str_list = [s.decode("utf-8") for s in strarr.tolist()]
        else:
            str_list = [str(s) for s in strarr.tolist()]
        values = "".join(str_list)
        if "\0" in values:
            warnings.warn(
                (
                    "Found embedded NUL (\\0) characters in string categories. "
                    "Arrow used to strip these characters, but they are now preserved."
                ),
                UserWarning,
            )
        return offsets.astype(np.int32), values

    # String index type
    name_offsets, name_values = npstr_to_arrow_strarr(cats.values)
    name_offsets, _ = _ensure_np_dtype(name_offsets, np.int32)
    joffsets = array_interface_dict(name_offsets)
    bvalues = name_values.encode("utf-8")

    ptr = ctypes.c_void_p.from_buffer(ctypes.c_char_p(bvalues)).value
    assert ptr is not None

    jvalues: ArrayInf = {
        "data": (ptr, True),
        "typestr": "|i1",
        "shape": (len(name_values),),
        "strides": None,
        "version": 3,
        "mask": None,
    }
    jnames: StringArray = {"offsets": joffsets, "values": jvalues}

    code_values = codes.values
    jcodes = array_interface_dict(code_values)

    buf = (
        name_offsets,
        name_values,
        bvalues,
        code_values,
    )  # store temporary values
    return jnames, jcodes, buf


def array_interface(data: np.ndarray) -> bytes:
    """Make array interface str."""
    interface = array_interface_dict(data)
    interface_str = bytes(json.dumps(interface), "utf-8")
    return interface_str


def check_cudf_meta(data: _CudaArrayLikeArg, field: str) -> None:
    "Make sure no missing value in meta data."
    if (
        "mask" in data.__cuda_array_interface__
        and data.__cuda_array_interface__["mask"] is not None
    ):
        raise ValueError(f"Missing value is not allowed for: {field}")


def _cudf_str_cat_inf(cats: DfCatAccessor) -> Tuple[CudaStringArray, Tuple]:
    """String category index path for :py:func:`cudf_cat_inf`."""
    import pylibcudf as plc  # pylint: disable=import-outside-toplevel

    # pylint: disable=protected-access
    plc_col = cats._column.to_pylibcudf()
    if plc_col.type().id() != plc.TypeId.STRING:
        raise TypeError(
            "Unexpected type for category index. It's neither numeric nor string."
        )
    # Categories should not have missing values nor a non-zero logical offset.
    assert plc_col.null_count() == 0
    assert plc_col.offset() == 0

    off_child = plc_col.children()[0]  # offsets
    assert off_child.type().id() == plc.TypeId.INT32, "Expected INT32 string offsets."

    # String category index in arrow format
    jdata: CudaArrayInf = _arrow_buf_inf(
        plc_col.data().__cuda_array_interface__["data"][0],
        "|i1",
        0,
        STREAM_PER_THREAD,
    )
    joffset: CudaArrayInf = _arrow_buf_inf(
        off_child.data().__cuda_array_interface__["data"][0],
        "<i4",
        off_child.size(),
        STREAM_PER_THREAD,
    )
    jnames: CudaStringArray = {"offsets": joffset, "values": jdata}
    # Keep `plc_col` alive: it owns the GPU buffers pointed to by `jdata` and
    # `joffset`.
    return jnames, (plc_col,)


def cudf_cat_inf(
    cats: DfCatAccessor, codes: "pd.Series"
) -> Tuple[Union[CudaArrayInf, CudaStringArray], ArrayInf, Tuple]:
    """Obtain the cuda array interface for cuDF categories."""
    cp = import_cupy()
    is_num_idx = cp.issubdtype(cats.dtype, cp.floating) or cp.issubdtype(
        cats.dtype, cp.integer
    )
    if is_num_idx:
        cats_ainf = cuda_array_interface_dict(cats)
        codes_ainf = cuda_array_interface_dict(codes)
        return cats_ainf, codes_ainf, (cats, codes)

    jnames, buf = _cudf_str_cat_inf(cats)
    jcodes = cuda_array_interface_dict(codes)
    return jnames, jcodes, buf


class Categories:
    """An internal storage class for categories returned by the DMatrix and the
    Booster. This class is designed to be opaque. It is intended to be used exclusively
    by XGBoost as an intermediate storage for re-coding categorical data.

    The categories are saved along with the booster object. As a result, users don't
    need to preserve this class for re-coding. Use the booster model IO instead if you
    want to preserve the categories in a stable format.

    .. versionadded:: 3.1.0

    .. warning::

        This class is internal.

    .. code-block:: python

        Xy = xgboost.QuantileDMatrix(X, y, enable_categorical=True)
        booster = xgboost.train({}, Xy)

        categories = booster.get_categories() # Get categories

        # Use categories as a reference for re-coding
        Xy_new = xgboost.QuantileDMatrix(
            X_new, y_new, feature_types=categories, enable_categorical=True, ref=Xy
        )

        # Categories will be part of the `model.json`.
        booster.save_model("model.json")

    """

    def __init__(
        self,
        handle: Tuple[ctypes.c_void_p, Callable[[], None]],
        arrow_arrays: Optional[ArrowCatList],
    ) -> None:
        # The handle type is a bundle of the handle and the free call. Otherwise, we
        # will have to import the `_lib` and the `_check_call` from the core module
        # inside the __del__ method to avoid cyclic model dependency.
        # Importing modules in __del__ can result in Python abort if __del__ is called
        # during exception handling (interpreter is shutting down).
        self._handle, self._free = handle
        self._arrow_arrays = arrow_arrays

    def to_arrow(self) -> ArrowCatList:
        """Get the categories in the dataset. The results are stored in a list of
        (feature name, arrow array) pairs, with one array for each categorical
        feature. If a feature is numerical, then the corresponding column in the list is
        None. A value error will be raised if this container was created without the
        `export_to_arrow` option.

        """
        if self._arrow_arrays is None:
            raise ValueError(
                "The `export_to_arrow` option of the `get_categories` method"
                " is required."
            )
        return self._arrow_arrays

    def empty(self) -> bool:
        """Returns True if there's no category."""
        return self._handle.value is None

    def get_handle(self) -> int:
        """Internal method for retrieving the handle."""
        assert self._handle.value
        return self._handle.value

    def __del__(self) -> None:
        if self._handle.value is None:
            return
        self._free()


def get_ref_categories(
    feature_types: Optional[Union[FeatureTypes, Categories]],
) -> Tuple[Optional[FeatureTypes], Optional[Categories]]:
    """Get the optional reference categories from the `feature_types`. This is used by
    various `DMatrix` where the `feature_types` is reused for specifying the reference
    categories.

    """
    if isinstance(feature_types, Categories):
        ref_categories = feature_types
        feature_types = None
    else:
        ref_categories = None
    return feature_types, ref_categories


# Type schema for storing JSON-encoded array interface
AifType: TypeAlias = List[
    Union[
        # numeric column
        Union[ArrayInf, CudaArrayInf],
        # categorical column
        Tuple[
            # (cuda) numeric index | (cuda) string index
            Union[ArrayInf, CudaArrayInf, StringArray, CudaStringArray],
            Union[ArrayInf, CudaArrayInf],  # codes
        ],
    ]
]


class TransformedDf(ABC):
    """Internal class for storing transformed dataframe.

    Parameters
    ----------
    ref_categories :
        Optional reference categories used for re-coding.

    aitfs :
        Array interface for each column.

    """

    def __init__(
        self,
        ref_categories: Optional[Categories],
        aitfs: AifType,
        temporary_buffers: List[Tuple],
    ) -> None:
        self.ref_categories = ref_categories
        if ref_categories is not None and ref_categories.get_handle() is not None:
            aif = ref_categories.get_handle()
            self.ref_aif: Optional[int] = aif
        else:
            self.ref_aif = None

        self.aitfs = aitfs
        self.temporary_buffers = temporary_buffers

    def array_interface(self) -> bytes:
        """Return a byte string for JSON encoded array interface."""
        if self.ref_categories is not None:
            ref_inf: dict = {"ref_categories": self.ref_aif, "columns": self.aitfs}
            inf = bytes(json.dumps(ref_inf), "utf-8")
        else:
            inf = bytes(json.dumps(self.aitfs), "utf-8")
        return inf

    @property
    @abstractmethod
    def shape(self) -> Tuple[int, int]:
        """Return the shape of the dataframe."""