from lit.llvm import llvm_config

if not config.mlir_include_integration_tests:
    config.unsupported = True


def configure_aarch64_mcr_cmd():
    mcr_cmd = "mlir-runner"

    # NOTE: If the SVE tests are disabled and the SME tests are enabled to run
    # under emulation, the SVE specific RUN lines in the SparseTensor tests
    # will run under emulation.
    if not (config.mlir_run_arm_sve_tests or config.mlir_run_arm_sme_tests):
        config.substitutions.append(("%mcr_aarch64_cmd", mcr_cmd))
        return

    if config.arm_emulator_executable:
        if config.arm_emulator_mlir_runner_executable:
            mcr_cmd = config.arm_emulator_mlir_runner_executable
        else:
            # Top-level LIT config adds llvm_tools_dir to PATH but this is lost
            # when running under an emulator. If the user didn't specify an
            # mlir-runner executable, use absolute path
            # %llvm_tools_dir/mlir-runner.
            mcr_cmd = llvm_config.use_llvm_tool(
                "mlir-runner",
                search_env="MLIR_RUNNER",
                required=True,
                search_paths=[config.mlir_tools_dir],
                use_installed=False,
            )

        # Run test in emulator (qemu or armie)
        emulation_cmd = (
            f"{config.arm_emulator_executable} {config.arm_emulator_options}"
        )

        mcr_cmd = f"{emulation_cmd} {mcr_cmd}"

    config.substitutions.append(("%mcr_aarch64_cmd", mcr_cmd))

def configure_aarch64_mlir_utils():
    # Use mlir runtime libs from the location passed by -DARM_EMULATOR_UTILS_LIB_DIR.
    if config.arm_emulator_utils_lib_dir:
      base_dir = config.arm_emulator_utils_lib_dir
    else:
      base_dir = config.mlir_lib_dir

    config.substitutions.append(("%native_mlir_runner_utils",
                                 base_dir + "/libmlir_runner_utils" + config.llvm_shlib_ext))
    config.substitutions.append(("%native_mlir_c_runner_utils",
                                 base_dir + "/libmlir_c_runner_utils" + config.llvm_shlib_ext))
    config.substitutions.append(("%native_mlir_arm_runner_utils",
                                 base_dir + "/libmlir_arm_runner_utils" + config.llvm_shlib_ext))
    # Use passed Arm SME ABI routines, if not present default to stubs.
    config.substitutions.append(("%native_arm_sme_abi_shlib", config.arm_sme_abi_routines_shlib or
                                 (base_dir + "/libmlir_arm_sme_abi_stubs" + config.llvm_shlib_ext)))

# AArch64 tests will run under emulation if configured at build time by the
# following CMake options:
#
# * ARM_EMULATOR_EXECUTABLE     - emulator to use.
# * ARM_EMULATOR_OPTIONS        - options for emulator.
# * ARM_EMULATOR_MLIR_RUNNER_EXECUTABLE - AArch64 native mlir-runner to
#                                             support cross-compilation
# * ARM_EMULATOR_UTILS_LIB_DIR - AArch64 native utilites library to support
#                                cross-compilation.

# Check for above options and set correct executables and runtime libs
configure_aarch64_mcr_cmd()
configure_aarch64_mlir_utils()