"""
Tests for the as_numba_type() machinery.
"""
import typing as py_typing


import unittest

from contextlib import ExitStack
from llvmlite import ir

from numba import jit
from numba.core import cgutils, types
from numba.core.datamodel.models import PrimitiveModel
from numba.core.errors import TypingError
from numba.core.extending import (register_model, type_callable, unbox,
                                  NativeValue)
from numba.core.types import Number
from numba.core.typing.typeof import typeof, typeof_impl
from numba.core.typing.asnumbatype import as_numba_type, AsNumbaTypeRegistry
from numba.core.utils import PYVERSION
from numba.experimental.jitclass import jitclass
from numba.tests.support import TestCase


class TestAsNumbaType(TestCase):

    int_nb_type = typeof(0)
    float_nb_type = typeof(0.0)
    complex_nb_type = typeof(complex(0))
    str_nb_type = typeof("numba")
    bool_nb_type = typeof(True)
    none_nb_type = typeof(None)

    def test_simple_types(self):
        self.assertEqual(as_numba_type(int), self.int_nb_type)
        self.assertEqual(as_numba_type(float), self.float_nb_type)
        self.assertEqual(as_numba_type(complex), self.complex_nb_type)
        self.assertEqual(as_numba_type(str), self.str_nb_type)
        self.assertEqual(as_numba_type(bool), self.bool_nb_type)
        self.assertEqual(as_numba_type(type(None)), self.none_nb_type)

    def test_numba_types(self):
        numba_types = [
            types.intp,
            types.boolean,
            types.ListType(types.float64),
            types.DictType(
                types.intp, types.Tuple([types.float32, types.float32])
            ),
        ]

        for ty in numba_types:
            self.assertEqual(as_numba_type(ty), ty)

    def test_single_containers(self):
        # Native container types
        self.assertEqual(
            as_numba_type(list[float]),
            types.ListType(self.float_nb_type),
        )
        self.assertEqual(
            as_numba_type(dict[float, str]),
            types.DictType(self.float_nb_type, self.str_nb_type),
        )
        self.assertEqual(
            as_numba_type(set[complex]),
            types.SetType(self.complex_nb_type),
        )
        self.assertEqual(
            as_numba_type(tuple[float, float]),
            types.Tuple([self.float_nb_type, self.float_nb_type]),
        )
        self.assertEqual(
            as_numba_type(tuple[float, complex]),
            types.Tuple([self.float_nb_type, self.complex_nb_type]),
        )

        # Corresponding deprecated aliases
        self.assertEqual(
            as_numba_type(py_typing.List[float]),
            types.ListType(self.float_nb_type),
        )
        self.assertEqual(
            as_numba_type(py_typing.Dict[float, str]),
            types.DictType(self.float_nb_type, self.str_nb_type),
        )
        self.assertEqual(
            as_numba_type(py_typing.Set[complex]),
            types.SetType(self.complex_nb_type),
        )
        self.assertEqual(
            as_numba_type(py_typing.Tuple[float, float]),
            types.Tuple([self.float_nb_type, self.float_nb_type]),
        )
        self.assertEqual(
            as_numba_type(py_typing.Tuple[float, complex]),
            types.Tuple([self.float_nb_type, self.complex_nb_type]),
        )

    def test_optional(self):
        self.assertEqual(
            as_numba_type(py_typing.Optional[float]),
            types.Optional(self.float_nb_type),
        )
        self.assertEqual(
            as_numba_type(py_typing.Union[str, None]),
            types.Optional(self.str_nb_type),
        )
        self.assertEqual(
            as_numba_type(py_typing.Union[None, bool]),
            types.Optional(self.bool_nb_type),
        )

        # Optional[x] is a special case of Union[x, None].  We raise a
        # TypingError if the right type is not NoneType.
        with self.assertRaises(TypingError) as raises:
            as_numba_type(py_typing.Union[int, float])
        self.assertIn("Cannot type Union that is not an Optional",
                      str(raises.exception))

    @unittest.skipUnless(PYVERSION >= (3, 14),
                         "Syntax only supported from 3.14 onwards")
    def test_optional_syntax_OR(self):
        self.assertEqual(as_numba_type(int | None),
                         types.Optional(self.int_nb_type))

    def test_nested_containers(self):
        self.assertEqual(
            as_numba_type(list[list[int]]),
            types.ListType(types.ListType(self.int_nb_type)),
        )
        self.assertEqual(
            as_numba_type(list[py_typing.List[int]]),
            types.ListType(types.ListType(self.int_nb_type)),
        )
        self.assertEqual(
            as_numba_type(py_typing.List[list[int]]),
            types.ListType(types.ListType(self.int_nb_type)),
        )
        self.assertEqual(
            as_numba_type(py_typing.List[py_typing.List[int]]),
            types.ListType(types.ListType(self.int_nb_type)),
        )
        self.assertEqual(
            as_numba_type(list[dict[float, bool]]),
            types.ListType(
                types.DictType(self.float_nb_type, self.bool_nb_type)
            ),
        )
        self.assertEqual(
            as_numba_type(py_typing.List[py_typing.Dict[float, bool]]),
            types.ListType(
                types.DictType(self.float_nb_type, self.bool_nb_type)
            ),
        )
        self.assertEqual(
            as_numba_type(
                py_typing.Set[py_typing.Tuple[py_typing.Optional[int], float]]),
            types.SetType(types.Tuple(
                [types.Optional(self.int_nb_type), self.float_nb_type])),
        )

    def test_jitclass_registers(self):

        @jitclass
        class MyInt:
            x: int

            def __init__(self, value):
                self.x = value

        self.assertEqual(as_numba_type(MyInt), MyInt.class_type.instance_type)

    def test_type_alias(self):
        Pair = py_typing.Tuple[int, int]
        ListOfPairs = py_typing.List[Pair]

        pair_nb_type = types.Tuple((self.int_nb_type, self.int_nb_type))
        self.assertEqual(as_numba_type(Pair), pair_nb_type)
        self.assertEqual(
            as_numba_type(ListOfPairs), types.ListType(pair_nb_type)
        )

    def test_overwrite_type(self):
        as_numba_type = AsNumbaTypeRegistry()
        self.assertEqual(as_numba_type(float), self.float_nb_type)
        as_numba_type.register(float, types.float32)
        self.assertEqual(as_numba_type(float), types.float32)
        self.assertNotEqual(as_numba_type(float), self.float_nb_type)

    def test_any_throws(self):
        Any = py_typing.Any

        any_types = [
            py_typing.Optional[Any],
            py_typing.List[Any],
            py_typing.Set[Any],
            py_typing.Dict[float, Any],
            py_typing.Dict[Any, float],
            py_typing.Tuple[int, Any],
        ]

        for bad_py_type in any_types:
            with self.assertRaises(TypingError) as raises:
                as_numba_type(bad_py_type)
            self.assertIn(
                "Cannot infer Numba type of Python type",
                str(raises.exception),
            )

    def test_non_homogenous_type_throws(self):
        # Non-generic types of native container are not strictly homogeneous,
        # therefore not supported. Also, tuples with inhomogeneous types are
        # excluded.
        native_no_args_types = [
            list,
            set,
            dict,
            tuple,
            tuple[int, list],
        ]

        for bad_py_type in native_no_args_types:
            with self.assertRaises(TypingError) as raises:
                as_numba_type(bad_py_type)
            self.assertIn(
                "Cannot infer Numba type of Python type",
                str(raises.exception),
            )

    def test_bad_union_throws(self):
        bad_unions = [
            py_typing.Union[str, int],
            py_typing.Union[int, type(None), py_typing.Tuple[bool, bool]],
        ]

        for bad_py_type in bad_unions:
            with self.assertRaises(TypingError) as raises:
                as_numba_type(bad_py_type)
            self.assertIn("Cannot type Union", str(raises.exception))

    def test_instance_check_usecase(self):
        # Demonstrates that registering the type class with as_numba_type
        # supports instance checks, at least for those subclasses supported by
        # the instance check (e.g. types.Number, etc.).
        #
        # To set up the test we need quite a lot of extension code to support
        # a new type independent of the existing types.

        # The Python class
        class bfloat16:
            def __init__(self, value):
                self._value = value

        # The Numba type class - we use a Number subclass both because it makes
        # sense for a new numeric type, and it's one of the types supported by
        # instance checks in Numba
        class _type_class_bfloat16(Number):
            def __init__(self):
                self.bitwidth = 16
                super().__init__(name="bfloat16")

        # The Numba type instance
        bfloat16_type = _type_class_bfloat16()

        # Register typing of the Python class for use as arguments and
        # constants
        @typeof_impl.register(bfloat16)
        def typeof_bfloat16(val, c):
            return bfloat16_type

        # A data model for the bfloat16 class. We don't need much actual
        # implementation so it doesn't matter too much what this is; a 16-bit
        # integer representation is sufficient.
        @register_model(_type_class_bfloat16)
        class _model_bfloat16(PrimitiveModel):
            def __init__(self, dmm, fe_type):
                be_type = ir.IntType(fe_type.bitwidth)
                super(_model_bfloat16, self).__init__(dmm, fe_type, be_type)

        # Ideally we pass in a value so we ensure that the instance check is
        # working with values dynamically passed in (preventing the whole check
        # being potentially optimized into a simple True or False). For this we
        # need an unboxing.
        @unbox(_type_class_bfloat16)
        def unbox_bfloat16(ty, obj, c):
            ll_type = c.context.get_argument_type(ty)
            val = cgutils.alloca_once(c.builder, ll_type)
            is_error_ptr = cgutils.alloca_once_value(c.builder,
                                                     cgutils.false_bit)

            with ExitStack() as stack:
                value_obj = c.pyapi.object_getattr_string(obj, "_value")

                with cgutils.early_exit_if_null(c.builder, stack, value_obj):
                    c.builder.store(cgutils.true_bit, is_error_ptr)

                value_native = c.unbox(types.uint16, value_obj)
                c.pyapi.decref(value_obj)

                with cgutils.early_exit_if(c.builder, stack,
                                           value_native.is_error):
                    c.builder.store(cgutils.true_bit, is_error_ptr)

                c.builder.store(value_native.value, val)

            return NativeValue(c.builder.load(val),
                               is_error=c.builder.load(is_error_ptr))

        # We never call bfloat16 to construct one inside a jitted function, but
        # we need this typing so that the type of the bfloat16 class can be
        # determined (it's the argument to the instance check).
        @type_callable(bfloat16)
        def type_bfloat16_ctor(context):
            # Note that the typer is never called in this test, because we
            # don't call bfloat16 - only the typing of it as a callable is
            # used.

            def typer(value):
                if isinstance(value, types.Integer):
                    return bfloat16_type

        # First we try the instance check without an as_numba_type
        # registration, to prove that it is necessary for instance checks to
        # work.
        @jit
        def instancecheck_no_ant_reg(x):
            return isinstance(x, bfloat16)

        # A "random" value to test with
        x_bf16 = bfloat16(0x4049)  # bfloat16(3.14)

        # Ensure the typing fails without the registration
        expected_message = r"Cannot infer Numba type of Python type.*bfloat16"
        with self.assertRaisesRegex(TypingError, expected_message):
            instancecheck_no_ant_reg(x_bf16)

        # Register the typing with as_numba_type so that we can expect instance
        # checks to work
        as_numba_type.register(bfloat16, bfloat16_type)

        # We define a new function to ensure all registrations are as-required
        @jit
        def instancecheck(x):
            return isinstance(x, bfloat16)

        # The instance check should be True for bfloat16 instances and False
        # otherwise.
        self.assertTrue(instancecheck(x_bf16))
        self.assertFalse(instancecheck(1))


if __name__ == '__main__':
    unittest.main()
