# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import inspect
import pprint
import os
from typing import Any, Dict, Callable

import warnings
from numba_cuda_mlir.numba_cuda import compiler, types, typing
from numba_cuda_mlir.numba_cuda.core import errors, funcdesc, sigutils
from numba_cuda_mlir.numba_cuda.core.compiler_lock import global_compiler_lock
from numba_cuda_mlir.numba_cuda.compiler import CompilerBase, DefaultPassBuilder
from numba_cuda_mlir.numba_cuda.flags import CUDAFlags
from numba_cuda_mlir.numba_cuda.core.options import ParallelOptions
from numba_cuda_mlir.numba_cuda.core.errors import (
    NumbaDebugInfoWarning,
)
from numba_cuda_mlir.numba_cuda.compiler import compile_result
from numba_cuda_mlir.descriptor import mlir_target
from numba_cuda_mlir.lowering_utilities.type_conversions import to_numba_type
from numba_cuda_mlir.numba_cuda.core.compiler_machinery import (
    FunctionPass,
    LoweringPass,
    PassManager,
    register_pass,
)
from numba_cuda_mlir.numba_cuda.core.typed_passes import (
    AnnotateTypes,
    IRLegalization,
    NopythonTypeInference,
)
from numba_cuda_mlir.numba_cuda.core.typed_passes import (
    fallback_context,
    type_inference_stage,
)

# numba-cuda forked numba.core.ir into numba.cuda.core.ir, creating a separate
# UndefinedType class and UNDEFINED singleton. The core SSA pass inserts
# numba.core.ir.UNDEFINED into phi nodes, but numba-cuda's typeinfer checks
# `iv is not ir.UNDEFINED` using its own copy — the identity check fails.
# Unify them so phi nodes with undefined incoming values are handled correctly.
import numba_cuda_mlir.numba_cuda.core.ir as _cuda_ir

from numba_cuda_mlir.numba_cuda.typing.typeof import Purpose, typeof
from numba_cuda_mlir.numba_cuda.core.untyped_passes import (
    LiteralUnroll,
    InlineInlinables,
)
from numba_cuda_mlir.numbair_transforms import (
    NumbaCudaMlirLiteralUnroll,
    NumbaCudaMlirInlineInlinables,
    PostInlineWholeFunctionPlanners,
)

import numba_cuda_mlir.mlir_lowering as lowering
import numba_cuda_mlir.mlir_optimization as opt
from numba_cuda_mlir._launch_config import (
    _LAUNCH_CONFIG_TRACKER_METADATA_KEY,
    _LAUNCH_CONFIG_TRACKER_OPTION,
)
from numba_cuda_mlir.decorators import mlir_jit
from numba_cuda_mlir.ast_transforms import apply_ast_transforms
from numba_cuda_mlir.errors import (
    InternalCompilerError,
    UserFacingInternalCompilerError,
    handle_lowering_error,
)


@register_pass(mutates_CFG=True, analysis_only=False)
class MLIRTypeInference(FunctionPass):
    _name = "mlir_type_inference"
    _raise_errors = True

    def __init__(self):
        FunctionPass.__init__(self)

    def run_pass(self, state):
        """
        Type inference and legalization
        """
        with fallback_context(
            state, 'Function "%s" failed type inference' % (state.func_id.func_name,)
        ):
            # Type inference
            typemap, return_type, calltypes, errs = type_inference_stage(
                state.typingctx,
                state.targetctx,
                state.func_ir,
                state.args,
                state.return_type,
                state.locals,
                raise_errors=self._raise_errors,
            )
            state.typemap = typemap
            # save errors in case of partial typing
            state.typing_errors = errs
            if self._raise_errors:
                state.return_type = return_type
            state.calltypes = calltypes

        def legalize_return_type(return_type, interp, targetctx):
            """
            Reject function object return types.
            """

            if isinstance(return_type, types.Function) or isinstance(return_type, types.Phantom):
                if self._raise_errors:
                    msg = "Can't return function object ({}) in nopython mode"
                    raise errors.NumbaTypeError(msg.format(return_type))

        with fallback_context(
            state, 'Function "%s" has invalid return type' % (state.func_id.func_name,)
        ):
            legalize_return_type(state.return_type, state.func_ir, state.targetctx)
        return True


class MLIRLibrary:
    def __init__(self, code, ptx="", mlir_str=""):
        self._code = code
        self._ptx = ptx
        self._mlir_str = mlir_str
        # Objects that need to be kept alive for the lifetime of the kernels
        # generated by this library, e.g., device arrays captured from global
        # scope whose pointers are baked into the generated PTX as constants.
        self.referenced_objects = {}

    def get_ir(self):
        return self._code

    def get_llvm_str(self):
        return self._mlir_str

    def get_asm_str(self, cc=None):
        return self._ptx


# "mutates_cfg" probably carries more semantic load than just whether the CFG
# has been mutated, so we set it to True here to ensure Numba doesn't hold some
# invalid assumptions
@register_pass(mutates_CFG=True, analysis_only=False)
class MLIRBackend(LoweringPass):
    _name = "mlir_backend"

    def __init__(self):
        LoweringPass.__init__(self)

    def run_pass(self, state):
        """
        Back-end: Packages lowering output in a compile result
        """
        signature = typing.signature(state.return_type, *state.args)
        targetctx = state.targetctx
        func_ir = state.func_ir
        typemap = state.typemap
        restype = state.return_type
        calltypes = state.calltypes
        flags = state.flags
        metadata = state.metadata
        fndesc = funcdesc.PythonFunctionDescriptor.from_specialized_function(
            func_ir,
            typemap,
            restype,
            calltypes,
            mangler=targetctx.mangler,
            inline=flags.forceinline,
            noalias=flags.noalias,
            abi_tags=[flags.get_mangle_string()],
        )

        code = (
            f"=============== Numba IR ===============\n\n"
            f"{func_ir.dump_to_string()}\n\n"
            f"=============== Numba IR Type Map ===============\n\n"
            f"{pprint.pformat(typemap)}\n\n"
        )

        target_options = metadata["targetoptions"]
        dump = target_options.get("dump_numba_ir", False) or target_options.get("dump", False)

        if dump:
            print(code)

        library = MLIRLibrary(code)

        lower = lowering.MLIRLower(targetctx, fndesc, func_ir, metadata=metadata)

        try:
            with targetctx.push_code_library(library):
                lower.lower_to_mlir()
                lower.verify_mlir_module()
        except UserFacingInternalCompilerError:
            raise
        except (NotImplementedError, InternalCompilerError) as exc:
            if os.environ.get("NUMBA_CUDA_MLIR_ICE_FULL_TB", "0").strip() == "1":
                raise exc
            handle_lowering_error(lower, func_ir)

        state.cr = compile_result(
            typing_context=state.typingctx,
            target_context=state.targetctx,
            typing_error=state.status.fail_reason,
            type_annotation=state.type_annotation,
            library=library,
            call_helper=None,
            signature=signature,
            fndesc=fndesc,
            metadata=metadata,
        )
        return True


def get_compiler_class(
    targetoptions: Dict[str, Any],
    launch_config_tracker=None,
):
    class MLIRCompiler(CompilerBase):
        def define_pipelines(self):
            dpb = DefaultPassBuilder
            pm = PassManager("mlir")

            untyped_passes = dpb.define_untyped_pipeline(self.state)
            # Replace passes with numba_cuda_mlir-specific versions:
            # - LiteralUnroll -> NumbaCudaMlirLiteralUnroll (supports both literal_unroll versions)
            # - InlineInlinables -> NumbaCudaMlirInlineInlinables (skips self-recursive functions)
            modified_passes = []
            for impl, desc in untyped_passes.passes:
                if impl is LiteralUnroll:
                    modified_passes.append((NumbaCudaMlirLiteralUnroll, desc))
                elif impl is InlineInlinables:
                    modified_passes.append((NumbaCudaMlirInlineInlinables, desc))
                    modified_passes.append(
                        (
                            PostInlineWholeFunctionPlanners,
                            "post-inline whole-function extension planners",
                        )
                    )
                else:
                    modified_passes.append((impl, desc))
            pm.passes.extend(modified_passes)

            typed_passes = dpb.define_typed_pipeline(self.state)

            def replace_type_inference_pass(implementation, description):
                if implementation is NopythonTypeInference:
                    return (MLIRTypeInference, "mlir frontend")
                else:
                    return (implementation, description)

            mlir_typed_passes = [
                replace_type_inference_pass(implementation, description)
                for implementation, description in typed_passes.passes
            ]

            pm.passes.extend(mlir_typed_passes)

            lowering_passes = self.define_mlir_lowering_pipeline(self.state)
            pm.passes.extend(lowering_passes.passes)

            pm.finalize()
            return [pm]

        def define_mlir_lowering_pipeline(self, state):
            pm = PassManager("mlir_lowering")
            # Numba legalization
            pm.add_pass(IRLegalization, "ensure IR is legal prior to lowering")
            pm.add_pass(AnnotateTypes, "annotate types")

            # Lowering to MLIR
            pm.add_pass(MLIRBackend, "Run MLIR backend")

            pm.finalize()
            return pm

        # Capture user targetoptions from the active stack into compiler metadata
        def __init__(
            self,
            typingctx,
            targetctx,
            library,
            args,
            return_type,
            flags,
            locals,
            call_conv=None,
            abi_info=None,
        ):
            # These are added to work with the vendored Numba/Numba-CUDA code.
            # They will eventually need implementing, as part of the
            # requirement to support them in line with Numba-CUDA, but for now
            # we accept these parameters if they are None so that callees can
            # correctly call it.
            if call_conv is not None:
                raise RuntimeError("MLIR compiler does not support call_conv")
            if abi_info is not None:
                raise RuntimeError("MLIRCompiler does not support abi_info")
            super().__init__(typingctx, targetctx, library, args, return_type, flags, locals)
            # Attach options early so all passes can see them via state.metadata
            self.state.metadata["targetoptions"] = targetoptions
            if launch_config_tracker is not None:
                self.state.metadata[_LAUNCH_CONFIG_TRACKER_METADATA_KEY] = launch_config_tracker

    return MLIRCompiler


@global_compiler_lock
def compile_mlir(pyfunc, return_type, args, targetoptions: Dict[str, Any]):
    from numba_cuda_mlir.install_registry import register_lowering
    from numba_cuda_mlir.tools import resolve_gpu_target

    register_lowering()
    launch_config_tracker = targetoptions.pop(_LAUNCH_CONFIG_TRACKER_OPTION, None)
    gpu_target = resolve_gpu_target(targetoptions)
    targetoptions["chip"] = gpu_target["chip"]

    # Apply AST transforms at compile time with argtypes available
    # This allows consteval to access argument types and target options
    pyfunc, transformed_source = apply_ast_transforms(pyfunc, targetoptions, args)

    typingctx = mlir_target.typing_context
    targetctx = mlir_target.target_context

    # Do not compile to native code
    flags = CUDAFlags()
    flags.no_compile = True
    flags.no_cpython_wrapper = True
    flags.no_cfunc_wrapper = True

    from numba_cuda_mlir.memory_management.config import is_nrt_enabled

    flags.nrt = is_nrt_enabled()

    flags.auto_parallel = ParallelOptions(targetoptions.get("parallel", False))
    flags.nvvm_options = {}
    is_lto = targetoptions.get("lto", False) or targetoptions.get("_compile_output") == "ltoir"
    if is_lto:
        flags.nvvm_options["gen-lto"] = None

    # Both debug and lineinfo turn on debug information in the compiled code,
    # but we keep them separate arguments in case we later want to overload
    # some other behavior on the debug flag. In particular, -opt=3 is not
    # supported with debug enabled, and enabling only lineinfo should not
    # affect the error model.
    flags.dbg_directives_only = targetoptions.get("lineinfo", False)
    flags.debuginfo = targetoptions.get("debug", False) or flags.dbg_directives_only

    if flags.debuginfo:
        try:
            inspect.getsource(pyfunc)
        except OSError:
            warnings.warn(
                NumbaDebugInfoWarning(
                    f"Could not find source for function: {pyfunc}. "
                    "Debug line information may be inaccurate."
                )
            )

    # We should offer the Python error model for easier debugging of numerical
    # errors - e.g. division by zero.
    flags.error_model = "python"

    # Run compilation pipeline
    from numba_cuda_mlir.numba_cuda.core.target_extension import target_override
    from numba_cuda_mlir.lowering_utilities import context as mlir_ctx
    from numba_cuda_mlir._mlir import ir as mlir_ir

    flags.lto = is_lto

    with target_override("numba_cuda_mlir"), mlir_ctx.get_context(), mlir_ir.Location.unknown():
        cres = compiler.compile_extra(
            typingctx=typingctx,
            targetctx=targetctx,
            func=pyfunc,
            args=args,
            return_type=return_type,
            flags=flags,
            locals={},
            pipeline_class=get_compiler_class(targetoptions, launch_config_tracker),
        )

    cres.metadata.pop(_LAUNCH_CONFIG_TRACKER_METADATA_KEY, None)
    cres.metadata["targetoptions"] = targetoptions
    cres.metadata["gpu_target"] = gpu_target
    cres.metadata["transformed_source"] = transformed_source

    return cres


@global_compiler_lock
def mlir_compiler_entry(
    pyfunc,
    func_args,
    targetoptions: Dict[str, Any],
    abi="c",
    abi_info=None,
    override_argtypes=None,
):
    """Compile a Python function to MLIR for a given set of argument types.

    :param pyfunc: The Python function to compile.
    :param func_args: The function arguments used to invoke pyfunc.
    :param debug: Whether to include debug info in the compiled code.
    :type debug: bool
    :param lineinfo: Whether to include a line mapping from the compiled code
                     to the source code. Usually this is used with optimized
                     code (since debug mode would automatically include this),
                     so we want debug info in the LLVM IR but only the line
                     mapping in the final output.
    :type lineinfo: bool
    :param abi: The ABI for a compiled function - either ``"numba"`` or
                ``"c"``. Note that the Numba ABI is not considered stable.
                The C ABI is only supported for device functions at present.
    :type abi: str
    :param abi_info: A dict of ABI-specific options. The ``"c"`` ABI supports
                     one option, ``"abi_name"``, for providing the wrapper
                     function's name. The ``"numba"`` ABI has no options.
    :type abi_info: dict
    :param override_argtypes: Optional tuple of Numba types to use instead of
                              inferring from func_args. Used for tuple argument
                              flattening where the actual args don't match the
                              original tuple types.
    :type override_argtypes: tuple or None
    :return: (code, resty): The compiled code and inferred return type
    :rtype: tuple
    """
    if abi not in ("numba", "c"):
        raise NotImplementedError(f"Unsupported ABI: {abi}")

    abi_info = abi_info or dict()
    default_values = list(pyfunc.__defaults__) if pyfunc.__defaults__ else []
    sig = inspect.signature(pyfunc)

    # construct a list of remaining arguments of default values
    func_args_remained = default_values[
        len(func_args) + len(default_values) - len(sig.parameters) :
    ]
    func_args_extended = func_args + func_args_remained

    # Check if function has type annotations - if so, use declared signature
    # instead of inferring from runtime arguments
    params = list(sig.parameters.values())
    argtypes = []

    for i, param in enumerate(params):
        if param.kind == inspect.Parameter.VAR_POSITIONAL:
            if override_argtypes is not None:
                vararg_types = list(override_argtypes[i:])
            else:
                vararg_types = [typeof(a, Purpose.argument) for a in func_args_extended[i:]]
            argtypes.append(types.StarArgTuple.from_types(vararg_types))
            break

        override_argtype = (
            override_argtypes[i]
            if override_argtypes is not None and i < len(override_argtypes)
            else None
        )
        if isinstance(override_argtype, types.Literal):
            argtypes.append(override_argtype)
            continue

        annotation = param.annotation

        # If there's a Numba type annotation, use it directly
        if annotation != inspect.Parameter.empty and isinstance(annotation, types.Type):
            argtypes.append(annotation)
        else:
            # Try to convert annotation to Numba type, otherwise infer from runtime
            if annotation != inspect.Parameter.empty:
                try:
                    argtypes.append(to_numba_type(annotation))
                except (TypeError, NotImplementedError):
                    # Use override type if provided, otherwise infer from arg
                    if override_argtype is not None:
                        argtypes.append(override_argtype)
                    else:
                        argtypes.append(typeof(func_args_extended[i], Purpose.argument))
            else:
                # Use override type if provided, otherwise infer from arg
                if override_argtype is not None:
                    argtypes.append(override_argtype)
                else:
                    argtypes.append(typeof(func_args_extended[i], Purpose.argument))

    args, return_type = sigutils.normalize_signature(tuple(argtypes))

    cres = compile_mlir(
        pyfunc=pyfunc,
        return_type=return_type,
        args=args,
        targetoptions=targetoptions,
    )

    opt.optimize(cres)

    return cres
