module MPI

const tag_name = "mpi"

const augment = raw"""
    # Can't use Preferences since we might be running this very early with a non-existing Manifest
    MPIPreferences_UUID = Base.UUID("3da0fdf6-3ccc-4f1b-acd9-58baa6c99267")
    const preferences = Base.get_preferences(MPIPreferences_UUID)

    const binary_to_abi = Dict(
        "MPIABI_jll" => "MPIABI",
        "MPICH_jll" => "MPICH",
        "MPICH_CUDA_jll" => "MPICH",
        "MPItrampoline_jll" => "MPItrampoline",
        "MicrosoftMPI_jll" => "MicrosoftMPI",
        "OpenMPI_jll" => "OpenMPI",
    )

    # Keep logic in sync with MPIPreferences.jl
    function augment_mpi!(platform)
        # Doesn't need to be `const` since we depend on MPIPreferences so we
        # invalidate the cache when it changes.

        # Determine which jll binary provides our MPI implementation.
        # Use MPIPreferences if set, otherwise use either MPICH or MicrosoftMPI (on Windows).
        # Note: MPIPreferences uses `Sys.iswindows()` without the `platform` argument.
        binary = get(preferences, "binary", Sys.iswindows(platform) ? "MicrosoftMPI_jll" : "MPICH_jll")

        # Convert the jll binary to an ABI
        if binary == "system"
            abi = get(preferences, "abi", nothing)
            abi === nothing && error("MPIPreferences: Inconsistent state detected, `binary` is set to system, but `abi` is not set.")
        else
            abi = get(binary_to_abi, binary, nothing)
            abi == nothing && error("Unknown MPI binary: $binary")
        end

        # Add an `mpi` entry to the platform if it's missing
        if !haskey(platform, "mpi")
            platform["mpi"] = abi
        end

        return platform
    end
    """

using BinaryBuilder, Pkg
using Base.BinaryPlatforms

mpi_abis = (
    ("MPIABI", PackageSpec(name="MPIABI_jll"), "0.1.3", p -> !Sys.iswindows(p)),
    ("MPICH", PackageSpec(name="MPICH_jll"), "5.0.1", p -> !Sys.iswindows(p)),
    ("MPItrampoline", PackageSpec(name="MPItrampoline_jll"), "5.5.6", p -> !Sys.iswindows(p) && !(libc(p) == "musl")),
    ("MicrosoftMPI", PackageSpec(name="MicrosoftMPI_jll"), "", Sys.iswindows),
    ("OpenMPI", PackageSpec(name="OpenMPI_jll"), "4.1.9, 5.0.11", p -> !Sys.iswindows(p) && !(arch(p) == "armv6l" && libc(p) == "glibc")),
)

"""
    augment_platforms(platforms; MPICH_compat = nothing, OpenMPI_compat = nothing, MicrosoftMPI_compat=nothing, MPItrampoline_compat=nothing)

This augments the platforms with different MPI versions. Compatibilities with different versions can be specified
"""
function augment_platforms(platforms;
                           MPIABI_compat = nothing,
                           MPICH_compat = nothing,
                           MPItrampoline_compat=nothing,
                           MicrosoftMPI_compat=nothing,
                           OpenMPI_compat = nothing)
    all_platforms = AbstractPlatform[]
    dependencies = []
    for (abi, pkg, compat, f) in mpi_abis

        # set specific versions of MPI packages
        if (abi=="MPIABI" && !isnothing(MPIABI_compat)) compat = MPIABI_compat; end
        if (abi=="MPICH" && !isnothing(MPICH_compat)) compat = MPICH_compat; end
        if (abi=="MPItrampoline" && !isnothing(MPItrampoline_compat)) compat = MPItrampoline_compat; end
        if (abi=="MicrosoftMPI" && !isnothing(MicrosoftMPI_compat)) compat = MicrosoftMPI_compat; end
        if (abi=="OpenMPI" && !isnothing(OpenMPI_compat)) compat = OpenMPI_compat; end

        pkg_platforms = deepcopy(filter(f, platforms))
        foreach(pkg_platforms) do p
            p[tag_name] = abi
        end
        append!(all_platforms, pkg_platforms)
        push!(dependencies, Dependency(pkg; compat, platforms=pkg_platforms))
    end
    # NOTE: packages using this platform tag must depend on MPIPreferences,
    #       otherwise they will not be invalidated when the Preference changes.
    push!(dependencies, RuntimeDependency(PackageSpec(name="MPIPreferences", uuid="3da0fdf6-3ccc-4f1b-acd9-58baa6c99267");
                                          compat="0.1", top_level=true))
    return all_platforms, dependencies
end

end
