load("@aspect_bazel_lib//lib:tar.bzl", "tar")
load("@bazel_skylib//rules:common_settings.bzl", "bool_flag")
load("@bazel_skylib//rules:write_file.bzl", "write_file")
load("@rules_zig//zig:defs.bzl", "zig_library")

RUNTIMES = {
    "cpu": True,
    "cuda": False,
    "rocm": False,
    "tpu": False,
    "neuron": False,
}

[
    bool_flag(
        name = runtime,
        build_setting_default = default,
    )
    for runtime, default in RUNTIMES.items()
]

[
    config_setting(
        name = "{}.enabled".format(runtime),
        flag_values = {":{}".format(runtime): "True"},
        visibility = ["//runtimes:__subpackages__"],
    )
    for runtime in RUNTIMES.keys()
]

write_file(
    name = "zml_txt",
    out = "zml.txt",
    content = ["ZML loves you <3"],
)

tar(
    name = "empty",
    srcs = [":zml_txt"],
    mtree = "auto",
)

filegroup(
    name = "layers",
    srcs = [":empty"] + select({
        ":cpu.enabled": ["//runtimes/cpu:layers"],
        "//conditions:default": [],
    }) + select({
        ":cuda.enabled": ["//runtimes/cuda:layers"],
        "//conditions:default": [],
    }) + select({
        ":neuron.enabled": ["//runtimes/neuron:layers"],
        "//conditions:default": [],
    }) + select({
        ":rocm.enabled": ["//runtimes/rocm:layers"],
        "//conditions:default": [],
    }) + select({
        ":tpu.enabled": ["//runtimes/tpu:layers"],
        "//conditions:default": [],
    }),
    visibility = ["//visibility:public"],
)

zig_library(
    name = "runtimes",
    main = "runtimes.zig",
    visibility = ["//visibility:public"],
    deps = [
        "//pjrt",
    ] + [
        "//runtimes/{}".format(runtime)
        for runtime in RUNTIMES.keys()
    ],
)