# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-2-Clause

from collections import namedtuple
from functools import singledispatch
import ctypes
import enum
import functools
import operator

import numpy as np
from numpy.random.bit_generator import BitGenerator
from numba.cuda import HAS_NUMBA
from numba.cuda import types
from numba.cuda.core import errors
from numba.cuda import utils
from numba.cuda.np import numpy_support

from cuda.core import Buffer
from cuda.core.utils import StridedMemoryView


# terminal color markup
_termcolor = errors.termcolor()


class Purpose(enum.Enum):
    # Value being typed is used as an argument
    argument = 1
    # Value being typed is used as a constant
    constant = 2


_TypeofContext = namedtuple("_TypeofContext", ("purpose",))


def typeof(val, purpose=Purpose.argument):
    """
    Get the Numba type of a Python value for the given purpose.
    """
    # Note the behaviour for Purpose.argument must match _typeof.c.
    c = _TypeofContext(purpose)
    ty = typeof_impl(val, c)
    if ty is None:
        msg = _termcolor.errmsg(f"Cannot determine Numba type of {type(val)}")
        raise ValueError(msg)
    return ty


@singledispatch
def typeof_impl(val, c):
    """
    Generic typeof() implementation.
    """
    tp = getattr(val, "_numba_type_", None)
    if tp is not None:
        return tp

    # Check for dlpack objects
    dlpack = getattr(val, "__dlpack__", None)
    if dlpack is not None:
        tp = _typeof_dlpack(dlpack, c)
        if tp is not None:
            return tp

    # Check for __cuda_array_interface__ objects (third-party device arrays)

    # Numba's own DeviceNDArray is handled above via _numba_type_.
    cai = getattr(val, "__cuda_array_interface__", None)
    if cai is not None:
        tp = _typeof_cuda_array_interface(cai, c)
        if tp is not None:
            return tp

    tp = _typeof_buffer(val, c)
    if tp is not None:
        return tp

    # cffi is handled here as it does not expose a public base class
    # for exported functions or CompiledFFI instances.
    from numba.cuda.typing import cffi_utils

    if cffi_utils.SUPPORTED:
        if cffi_utils.is_cffi_func(val):
            return cffi_utils.make_function_type(val)
        if cffi_utils.is_ffi_instance(val):
            return types.ffi

    if HAS_NUMBA:
        # Fallback to Numba's typeof_impl for third-party registrations
        from numba.core.typing.typeof import typeof_impl as core_typeof_impl

        tp = core_typeof_impl(val, c)
        if tp is not None:
            return tp

    return None


def _typeof_buffer(val, c):
    from numba.cuda.typing import bufproto

    try:
        m = memoryview(val)
    except TypeError:
        return
    # Object has the buffer protocol
    try:
        dtype = bufproto.decode_pep3118_format(m.format, m.itemsize)
    except ValueError:
        return
    type_class = bufproto.get_type_class(type(val))
    layout = bufproto.infer_layout(m)
    return type_class(dtype, m.ndim, layout=layout, readonly=m.readonly)


@typeof_impl.register(ctypes._CFuncPtr)
def _typeof_ctypes_function(val, c):
    from .ctypes_utils import is_ctypes_funcptr, make_function_type

    if is_ctypes_funcptr(val):
        return make_function_type(val)


@typeof_impl.register(type)
def _typeof_type(val, c):
    """
    Type various specific Python types.
    """
    if issubclass(val, BaseException):
        return types.ExceptionClass(val)
    if issubclass(val, tuple) and hasattr(val, "_asdict"):
        return types.NamedTupleClass(val)

    if issubclass(val, np.generic):
        return types.NumberClass(numpy_support.from_dtype(val))

    if issubclass(val, types.Type):
        return types.TypeRef(val)


@typeof_impl.register(bool)
def _typeof_bool(val, c):
    return types.boolean


@typeof_impl.register(float)
def _typeof_float(val, c):
    return types.float64


@typeof_impl.register(complex)
def _typeof_complex(val, c):
    return types.complex128


@typeof_impl.register(int)
def _typeof_int(val, c):
    # As in _typeof.c
    nbits = utils.bit_length(val)
    if nbits < 32:
        typ = types.intp
    elif nbits < 64:
        typ = types.int64
    elif nbits == 64 and val >= 0:
        typ = types.uint64
    else:
        raise ValueError("Int value is too large: %s" % val)
    return typ


@typeof_impl.register(np.generic)
def _typeof_numpy_scalar(val, c):
    try:
        return numpy_support.map_arrayscalar_type(val)
    except errors.NumbaNotImplementedError:
        pass
    except NotImplementedError:
        pass


@typeof_impl.register(str)
def _typeof_str(val, c):
    return types.string


@typeof_impl.register(type((lambda a: a).__code__))
def _typeof_code(val, c):
    return types.code_type


@typeof_impl.register(type(None))
def _typeof_none(val, c):
    return types.none


@typeof_impl.register(type(Ellipsis))
def _typeof_ellipsis(val, c):
    return types.ellipsis


@typeof_impl.register(tuple)
def _typeof_tuple(val, c):
    tys = [typeof_impl(v, c) for v in val]
    if any(ty is None for ty in tys):
        return
    return types.BaseTuple.from_types(tys, type(val))


@typeof_impl.register(list)
def _typeof_list(val, c):
    if len(val) == 0:
        raise ValueError("Cannot type empty list")
    ty = typeof_impl(val[0], c)
    if ty is None:
        raise ValueError(f"Cannot type list element type {type(val[0])}")
    return types.List(ty, reflected=True)


@typeof_impl.register(set)
def _typeof_set(val, c):
    if len(val) == 0:
        raise ValueError("Cannot type empty set")
    item = next(iter(val))
    ty = typeof_impl(item, c)
    if ty is None:
        raise ValueError(f"Cannot type set element type {type(item)}")
    return types.Set(ty, reflected=True)


@typeof_impl.register(slice)
def _typeof_slice(val, c):
    return types.slice2_type if val.step in (None, 1) else types.slice3_type


@typeof_impl.register(enum.Enum)
@typeof_impl.register(enum.IntEnum)
def _typeof_enum(val, c):
    clsty = typeof_impl(type(val), c)
    return clsty.member_type


@typeof_impl.register(enum.EnumMeta)
def _typeof_enum_class(val, c):
    cls = val
    members = list(cls.__members__.values())
    if len(members) == 0:
        raise ValueError("Cannot type enum with no members")
    dtypes = {typeof_impl(mem.value, c) for mem in members}
    if len(dtypes) > 1:
        raise ValueError(
            "Cannot type heterogeneous enum: "
            "got value types %s" % ", ".join(sorted(str(ty) for ty in dtypes))
        )
    if issubclass(val, enum.IntEnum):
        typecls = types.IntEnumClass
    else:
        typecls = types.EnumClass
    return typecls(cls, dtypes.pop())


@typeof_impl.register(np.dtype)
def _typeof_dtype(val, c):
    tp = numpy_support.from_dtype(val)
    return types.DType(tp)


@typeof_impl.register(np.ndarray)
def _typeof_ndarray(val, c):
    if isinstance(val, np.ma.MaskedArray):
        msg = "Unsupported array type: numpy.ma.MaskedArray."
        raise errors.NumbaTypeError(msg)
    try:
        dtype = numpy_support.from_dtype(val.dtype)
    except errors.NumbaNotImplementedError:
        raise errors.NumbaValueError(f"Unsupported array dtype: {val.dtype}")
    layout = numpy_support.map_layout(val)
    readonly = not val.flags.writeable
    return types.Array(dtype, val.ndim, layout, readonly=readonly)


@typeof_impl.register(types.NumberClass)
def _typeof_number_class(val, c):
    return val


@typeof_impl.register(types.Literal)
def _typeof_literal(val, c):
    return val


@typeof_impl.register(types.TypeRef)
def _typeof_typeref(val, c):
    return val


@typeof_impl.register(types.Type)
def _typeof_nb_type(val, c):
    if isinstance(val, types.BaseFunction):
        return val
    elif isinstance(val, (types.Number, types.Boolean)):
        return types.NumberClass(val)
    else:
        return types.TypeRef(val)


@typeof_impl.register(BitGenerator)
def typeof_numpy_random_bitgen(val, c):
    return types.NumPyRandomBitGeneratorType(val)


@typeof_impl.register(np.random.Generator)
def typeof_random_generator(val, c):
    return types.NumPyRandomGeneratorType(val)


@typeof_impl.register(np.polynomial.polynomial.Polynomial)
def typeof_numpy_polynomial(val, c):
    coef = typeof(val.coef)
    domain = typeof(val.domain)
    window = typeof(val.window)
    return types.PolynomialType(coef, domain, window)


@functools.lru_cache
def _typeof_cuda_array_interface_cached(
    *, dtype, shape, strides, readonly: bool
):
    # Determine layout
    ndim = len(shape)
    if not ndim:
        layout = "C"
    elif strides is None:
        layout = "C"
    else:
        itemsize = dtype.bitwidth
        # Quick rejection: C-contiguous has strides[-1] == itemsize,
        # F-contiguous has strides[0] == itemsize. If neither, it's "A".
        if strides[-1] == itemsize:
            c_strides = numpy_support.strides_from_shape(
                shape=shape,
                itemsize=itemsize,
                c_contiguous=True,
                f_contiguous=False,
            )
            layout = "C" if all(map(operator.eq, strides, c_strides)) else "A"
        elif strides[0] == itemsize:
            f_strides = numpy_support.strides_from_shape(
                shape=shape,
                itemsize=itemsize,
                c_contiguous=False,
                f_contiguous=True,
            )
            layout = "F" if all(map(operator.eq, strides, f_strides)) else "A"
        else:
            layout = "A"

    return types.Array(dtype, ndim, layout, readonly=readonly)


def _typeof_dlpack(val, c):
    obj = getattr(val, "__self__", None)
    if obj is not None:
        try:
            smv = StridedMemoryView.from_dlpack(obj, stream_ptr=-1)
        except BufferError:
            return

        smv_layout = smv._layout
        layout = (
            "C"
            if smv_layout.is_contiguous_c
            else "F"
            if smv_layout.is_contiguous_f
            else "A"
        )
        return types.Array(
            dtype=numpy_support.from_dtype(smv.dtype),
            ndim=len(smv.shape),
            layout=layout,
            readonly=smv.readonly,
        )


@functools.lru_cache
def _numba_dtype_from_str(typestr):
    return numpy_support.from_dtype(np.dtype(typestr))


def _typeof_cuda_array_interface(val, c):
    """
    Determine the type of a __cuda_array_interface__ object.

    This handles third-party device arrays that implement the CUDA
    Array Interface. These are typed as regular Array types, with lowering
    handled in numba.cuda.np.arrayobj.
    """
    dtype = _numba_dtype_from_str(val["typestr"])
    shape = val["shape"]
    strides = val.get("strides")
    _, readonly = val["data"]
    return _typeof_cuda_array_interface_cached(
        dtype=dtype,
        shape=shape,
        strides=strides,
        readonly=readonly,
    )


@typeof_impl.register(Buffer)
def typeof_buffer(val, c):
    return types.Array(
        dtype=_numba_dtype_from_str("uint8"),
        ndim=1,
        layout="C",
        readonly=False,
    )


@typeof_impl.register(StridedMemoryView)
def typeof_strided_memory_view(val, c):
    raw_layout = val._layout
    if raw_layout.is_contiguous_c:
        layout = "C"
    elif raw_layout.is_contiguous_f:
        layout = "F"
    elif raw_layout.is_contiguous_any:
        layout = "A"
    else:
        raise ValueError(
            "Unsupported StridedMemoryView layout; must be contiguous"
        )
    return types.Array(
        dtype=numpy_support.from_dtype(val.dtype),
        ndim=len(val.shape),
        layout=layout,
        readonly=val.readonly,
    )
