import inspect
import typing as py_typing

from numba.core.typing.typeof import typeof
from numba.core import errors, types
from numba.core.utils import PYVERSION


class AsNumbaTypeRegistry:
    """
    A registry for Python types. It stores a lookup table for simple cases
    (e.g. ``int``) and a list of functions for more complicated cases (e.g.
    generics like ``List[int]``).

    Python types are used in Python type annotations, and in instance checks.
    Therefore, this registry supports determining the Numba type of Python type
    annotations at compile time, along with determining the type of classinfo
    arguments to ``isinstance()``.

    This registry is not used dynamically on instances at runtime; to check the
    type of an object at runtime, use ``numba.typeof``.
    """

    def __init__(self):
        self.lookup = {
            type(example): typeof(example)
            for example in [
                0,
                0.0,
                complex(0),
                "numba",
                True,
                None,
            ]
        }

        self.functions = [self._builtin_infer, self._numba_type_infer]

    def _numba_type_infer(self, py_type):
        if isinstance(py_type, types.Type):
            return py_type

    def _builtin_infer(self, py_type):
        if PYVERSION in ((3, 14), ):
            # As of 3.14 the typing module has been updated to return a
            # different type when calling: `typing.Optional[X]`.
            #
            # On 3.14:
            #
            # >>> type(typing.Optional[float])
            # <class 'typing.Union'>
            #
            #
            # On 3.13 (and presumably below):
            #
            # >>> type(typing._UnionGenericAlias)
            # <class 'typing._UnionGenericAlias'>
            #
            #
            # The previous implementation of this predicate used
            # `_GenericAlias`, which was possible because `_UnionGenericAlias`
            # is a subclass of `_GenericAlias`...
            #
            # >>> issubclass(typing._UnionGenericAlias, typing._GenericAlias)
            # True
            #
            # However, other types, such as `typing.List[float]` remain as
            # `typing._GenericAlias`, so that must be keept.
            #
            # Additionally, using the recommend e.g. `tuple[int, float]`,
            # creates a `typing.GenericAlias`, so that must be included here
            # too.
            #
            if not isinstance(py_type, (py_typing.Union,
                                        py_typing.GenericAlias,
                                        py_typing._GenericAlias)):
                return
        elif PYVERSION in ((3, 10), (3, 11), (3, 12), (3, 13)):
            # Subscripting a class, e.g. `tuple[int, float]`, creates a
            # `typing.GenericAlias`. Meanwhile, using deprecated aliases such
            # as `typing.Tuple[int, float]` creates a `typing._GenericAlias`.
            if not isinstance(py_type, (py_typing.GenericAlias,
                                        py_typing._GenericAlias)):
                return
        else:
            raise NotImplementedError(PYVERSION)

        origin = py_typing.get_origin(py_type)
        args = py_typing.get_args(py_type)

        if origin is py_typing.Union:
            if len(args) != 2:
                raise errors.TypingError(
                    "Cannot type Union of more than two types. "
                    f"Attempted to unify '{len(args)}' types.")

            (arg_1_py, arg_2_py) = args

            if arg_2_py is type(None): # noqa: E721
                return types.Optional(self.infer(arg_1_py))
            elif arg_1_py is type(None): # noqa: E721
                return types.Optional(self.infer(arg_2_py))
            else:
                raise errors.TypingError(
                    "Cannot type Union that is not an Optional "
                    f"(neither type type {arg_2_py} is not NoneType")

        if origin is list:
            (element_py,) = args
            return types.ListType(self.infer(element_py))

        if origin is dict:
            key_py, value_py = args
            return types.DictType(self.infer(key_py), self.infer(value_py))

        if origin is set:
            (element_py,) = args
            return types.SetType(self.infer(element_py))

        if origin is tuple:
            tys = tuple(map(self.infer, args))
            return types.BaseTuple.from_types(tys)

    def register(self, func_or_py_type, numba_type=None):
        """
        Add support for new Python types (e.g. user-defined JitClasses) to the
        registry. For a simple pair of a Python type and a Numba type, this can
        be called as a function ``register(py_type, numba_type)``. If more
        complex logic is required (e.g. for generic types), ``register`` can be
        used as a decorator for a function that takes a Python type as input
        and returns a Numba type or ``None``.
        """
        if numba_type is not None:
            # register used with a specific (py_type, numba_type) pair.
            assert isinstance(numba_type, types.Type)
            self.lookup[func_or_py_type] = numba_type
        else:
            # register used as a decorator.
            assert inspect.isfunction(func_or_py_type)
            self.functions.append(func_or_py_type)

    def try_infer(self, py_type):
        """
        Try to determine the Numba type of a given Python type. We first
        consider the lookup dictionary. If ``py_type`` is not there, we iterate
        through the registered functions until one returns a Numba type.  If
        type inference fails, return ``None``.
        """
        result = self.lookup.get(py_type, None)

        for func in self.functions:
            if result is not None:
                break
            result = func(py_type)

        if result is not None and not isinstance(result, types.Type):
            raise errors.TypingError(
                f"as_numba_type should return a Numba type, got {result}"
            )
        return result

    def infer(self, py_type):
        result = self.try_infer(py_type)
        if result is None:
            raise errors.TypingError(
                f"Cannot infer Numba type of Python type {py_type}"
            )
        return result

    def __call__(self, py_type):
        return self.infer(py_type)


as_numba_type = AsNumbaTypeRegistry()
