# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE

# v2: replace with src/awkward/_v2/_connect/numba/__init__.py.

import awkward as ak

checked_version = False


def register_and_check():
    global checked_version
    try:
        import numba
    except ModuleNotFoundError:
        raise ModuleNotFoundError(
            """install the 'numba' package with:

    pip install numba --upgrade

or

    conda install numba"""
        ) from None
    else:
        if not checked_version and ak._v2._util.parse_version(
            numba.__version__
        ) < ak._v2._util.parse_version("0.50"):
            raise ImportError(
                "Awkward Array can only work with numba 0.50 or later "
                "(you have version {})".format(numba.__version__)
            )
        checked_version = True
        register()


def register():
    import numba
    import awkward._connect._numba.arrayview
    import awkward._connect._numba.layout
    import awkward._connect._numba.builder

    if hasattr(ak.numba, "ArrayViewType"):
        return

    n = ak.numba
    n.ArrayViewType = awkward._connect._numba.arrayview.ArrayViewType
    n.ArrayViewModel = awkward._connect._numba.arrayview.ArrayViewModel
    n.RecordViewType = awkward._connect._numba.arrayview.RecordViewType
    n.RecordViewModel = awkward._connect._numba.arrayview.RecordViewModel
    n.ContentType = awkward._connect._numba.layout.ContentType
    n.NumpyArrayType = awkward._connect._numba.layout.NumpyArrayType
    n.RegularArrayType = awkward._connect._numba.layout.RegularArrayType
    n.ListArrayType = awkward._connect._numba.layout.ListArrayType
    n.IndexedArrayType = awkward._connect._numba.layout.IndexedArrayType
    n.IndexedOptionArrayType = awkward._connect._numba.layout.IndexedOptionArrayType
    n.ByteMaskedArrayType = awkward._connect._numba.layout.ByteMaskedArrayType
    n.BitMaskedArrayType = awkward._connect._numba.layout.BitMaskedArrayType
    n.UnmaskedArrayType = awkward._connect._numba.layout.UnmaskedArrayType
    n.RecordArrayType = awkward._connect._numba.layout.RecordArrayType
    n.UnionArrayType = awkward._connect._numba.layout.UnionArrayType
    n.ArrayBuilderType = awkward._connect._numba.builder.ArrayBuilderType
    n.ArrayBuilderModel = awkward._connect._numba.builder.ArrayBuilderModel

    @numba.extending.typeof_impl.register(ak.highlevel.Array)
    def typeof_Array(obj, c):
        return obj.numba_type

    @numba.extending.typeof_impl.register(ak.highlevel.Record)
    def typeof_Record(obj, c):
        return obj.numba_type

    @numba.extending.typeof_impl.register(ak.highlevel.ArrayBuilder)
    def typeof_ArrayBuilder(obj, c):
        return obj.numba_type


def repr_behavior(behavior):
    return repr(behavior)


def castint(context, builder, fromtype, totype, val):
    import numba
    import llvmlite.ir.types

    if isinstance(fromtype, llvmlite.ir.types.IntType):
        if fromtype.width == 8:
            fromtype = numba.int8
        elif fromtype.width == 16:
            fromtype = numba.int16
        elif fromtype.width == 32:
            fromtype = numba.int32
        elif fromtype.width == 64:
            fromtype = numba.int64
    if not isinstance(fromtype, numba.types.Integer):
        raise AssertionError(
            f"unrecognized integer type: {fromtype!r}"
            + ak._util.exception_suffix(__file__)
        )

    if fromtype.bitwidth < totype.bitwidth:
        if fromtype.signed:
            return builder.sext(val, context.get_value_type(totype))
        else:
            return builder.zext(val, context.get_value_type(totype))
    elif fromtype.bitwidth > totype.bitwidth:
        return builder.trunc(val, context.get_value_type(totype))
    else:
        return val
