load("@bazel_lib//lib:copy_to_directory.bzl", "copy_to_directory")
load("@bazel_lib//lib:transitions.bzl", "platform_transition_filegroup")
load("@bazel_lib//tools/release:hashes.bzl", "hashes")
load("@bazel_skylib//rules:run_binary.bzl", "run_binary")
load("@rules_oci//oci:defs.bzl", "oci_image", "oci_load")
load("@rules_shell//shell:sh_binary.bzl", "sh_binary")
load("@rules_zig//zig:defs.bzl", "zig_binary", "zig_library")
load("@tar.bzl//tar:mtree.bzl", "mtree_spec")
load("@tar.bzl//tar:tar.bzl", "tar")

# zig libraries

# smi

run_binary(
    name = "libnrt.stripped",
    srcs = ["@aws-neuronx-runtime-lib//:libnrt.patchelf"],
    outs = ["libnrt.so.1"],
    args = [
        "--strip-debug",
        "-o",
        "$(location libnrt.so.1)",
        "$(location @aws-neuronx-runtime-lib//:libnrt.patchelf)",
    ],
    tool = "@llvm//tools:llvm-strip",
)

run_binary(
    name = "libncfw.stripped",
    srcs = ["@aws-neuronx-runtime-lib//:libncfw.patchelf"],
    outs = ["libncfw.so.2"],
    args = [
        "--strip-debug",
        "-o",
        "$(location libncfw.so.2)",
        "$(location @aws-neuronx-runtime-lib//:libncfw.patchelf)",
    ],
    tool = "@llvm//tools:llvm-strip",
)

copy_to_directory(
    name = "sandbox",
    srcs = select({
        "@platforms//os:macos": [],
        "@llvm//platforms/config:linux_x86_64": [
            ":libncfw.stripped",
            ":libnrt.stripped",
            "//platforms/rocm:zmlxrocm",
            "@amd-smi-lib//:libamd_smi",
            "@libdrm-amdgpu-amdgpu1",
            "@libdrm2-amdgpu//:libdrm2-amdgpu",
            "@libpjrt_rocm//:amdgpu_ids",
        ],
        "@llvm//platforms/config:linux_aarch64": [],
    }),
    include_external_repositories = ["**"],
    replace_prefixes = {
        "opt/amdgpu/lib/x86_64-linux-gnu": "lib",
        "amdgpu.ids": "share/libdrm/amdgpu.ids",
        "libdrm-amdgpu-amdgpu1": "lib",
        "platforms/rocm": "lib",
        "libnrt.so.1": "lib/libnrt.so.1",
        "libncfw.so.2": "lib/libncfw.so.2",
    },
)

run_binary(
    name = "zml-smi.stripped",
    srcs = [":zml-smi"],
    outs = ["zml-smi.stripped"],
    args = [
        "--strip-all",
        "-o",
        "$(location zml-smi.stripped)",
        "$(location :zml-smi)",
    ],
    tool = "@llvm//tools:llvm-strip",
)

zig_binary(
    name = "zml-smi",
    visibility = ["//visibility:public"],
    data = [":sandbox"],
    main = "main.zig",
    deps = [
        "//bin/zml-smi/collector",
        "//bin/zml-smi/csv",
        "//bin/zml-smi/json",
        "//bin/zml-smi/api",
        "//bin/zml-smi/prometheus",
        "//bin/zml-smi/utils:sysfs",
        "//bin/zml-smi/info",
        "//bin/zml-smi/tui",
        "//stdx",
        "@libvaxis//:vaxis",
        "@rules_zig//zig/runfiles",
    ] + select({
        "@platforms//os:macos": ["//bin/zml-smi/platforms/macos"],
        "@platforms//os:linux": ["//bin/zml-smi/platforms/linux"],
    }) + select({
        "@llvm//platforms/config:linux_x86_64": [
            "//bin/zml-smi/platforms/amdsmi",
            "//bin/zml-smi/platforms/oneapi",
            "//bin/zml-smi/platforms/neuron",
            "//bin/zml-smi/platforms/nvml",
            "//bin/zml-smi/platforms/tpu",
            "//platforms/neuron:libnrt_headers",
        ],
        "@llvm//platforms/config:linux_aarch64": [
            "//bin/zml-smi/platforms/nvml",
        ],
        "//conditions:default": [],
    }),
)

# docker image

mtree_spec(
    name = "mtree_spec",
    srcs = [
        ":zml-smi",
        ":zml-smi.stripped",
    ],
    tags = ["manual"],
)

# Use this genrule until mtree_mutate is fixed
genrule(
    name = "mtree",
    srcs = [":mtree_spec"],
    outs = ["mtree.txt"],
    cmd = """sed -e '/type=dir/d' -e '/zml-smi\\/zml-smi /d' -e 's#^bin/##g' -e 's#zml-smi\\.stripped#zml-smi#' -e '/stamp.c/d' $< > $@""",
    tags = ["manual"],
)

# mtree_mutate(
#     name = "mtree",
#     mtree = ":mtree_spec",
#     strip_prefix = "bin",
# )

# Release archives, one per platform, correctly named for distribution
# Build one:  bazel build --config=release //bin/zml-smi:zml-smi-linux-amd64
# Build all:  bazel build --config=release //bin/zml-smi:archives

_RELEASE_PLATFORMS = [
    ("linux", "amd64"),
    ("linux", "arm64"),
    ("macos", "arm64"),
]

[tar(
    name = "archive_" + os + "_" + arch,
    srcs = [
        ":zml-smi",
        ":zml-smi.stripped",
    ],
    out = "zml-smi-" + os + "-" + arch + ".tar.zst",
    args = [
        "--options",
        ",".join([
            "zstd:compression-level=22",
            "zstd:threads=16",
        ]),
    ],
    compress = "zstd",
    mtree = ":mtree",
    tags = ["manual"],
) for os, arch in _RELEASE_PLATFORMS]

[platform_transition_filegroup(
    name = "zml-smi-" + os + "-" + arch,
    srcs = [":archive_" + os + "_" + arch],
    target_platform = "//platforms:" + os + "_" + arch,
    tags = ["manual"],
) for os, arch in _RELEASE_PLATFORMS]

[hashes(
    name = "zml-smi-" + os + "-" + arch + ".sha256",
    src = ":zml-smi-" + os + "-" + arch,
    tags = ["manual"],
) for os, arch in _RELEASE_PLATFORMS]

sh_binary(
    name = "release",
    srcs = ["scripts/release.sh"],
    tags = ["manual"],
)

filegroup(
    name = "archives",
    srcs = [":zml-smi-" + os + "-" + arch for os, arch in _RELEASE_PLATFORMS] +
           [":zml-smi-" + os + "-" + arch + ".sha256" for os, arch in _RELEASE_PLATFORMS],
    tags = ["manual"],
)

oci_image(
    name = "image_",
    base = "@distroless_cc_debian12",
    entrypoint = ["/zml-smi/zml-smi"],
    target_compatible_with = [
        "@platforms//os:linux",
        "@platforms//cpu:x86_64",
    ],
    tars = [
        ":archive_linux_amd64",
    ],
    tags = ["manual"],
)

platform_transition_filegroup(
    name = "image",
    srcs = [":image_"],
    tags = ["manual"],
    target_platform = "//platforms:linux_amd64",
)

oci_load(
    name = "load",
    image = ":image",
    repo_tags = ["zmlai/zml-smi:latest"],
    tags = ["manual"],
)