load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("@rules_python//python/uv:lock.bzl", uv_lock = "lock")
load("@rules_zig//zig:defs.bzl", "zig_binary", "zig_library", "zig_shared_library")
load("@zml//bazel:cc_import.bzl", "cc_import")
load("@zml//bazel:runfiles.bzl", "runfiles_to_default")

exports_files([
    "neuronx_cc_shim.py",
    "nki_kernel_compiler.py",
])

# A proxy PJRT Plugin that loads the Neuron PJRT Plugin
# and returns the instance from nested GetPjrtApi.
#
# The proxy also initializes the Neuron runtime environment and exposes the
# sandboxed compiler binaries to libneuronxla 3.0.
zig_shared_library(
    name = "libpjrt_neuron",
    srcs = ["neuron.zig"],
    main = "libpjrt_neuron.zig",
    target_compatible_with = ["@platforms//os:linux"],
    visibility = ["@libpjrt_neuron//:__subpackages__"],
    zigopts = ["-fno-stack-check"],
    deps = [
        ":libnrt_headers",
        ":neuron_nrt",
        ":zmlxneuron",
        "//bazel",
        "//pjrt",
        "//stdx",
        "@rules_zig//zig/runfiles",
        "@xla//xla/pjrt/c:pjrt_c_api_hdrs",
    ],
)

zig_binary(
    name = "neuronx-cc",
    data = ["@neuron_py_deps//neuronx_cc"],
    linkopts = ["-Wl,-rpath,$$ORIGIN/../lib"],
    main = "neuronxcc_compiler_launcher.zig",
    tags = ["manual"],
    target_compatible_with = ["@platforms//os:linux"],
    deps = [
        ":libpython",
        ":python_launcher",
    ],
)

runfiles_to_default(
    name = "neuronx-cc_files",
    target_compatible_with = ["@platforms//os:linux"],
    visibility = ["@libpjrt_neuron//:__subpackages__"],
    deps = [":neuronx-cc"],
)

runfiles_to_default(
    name = "nki-cc_files",
    target_compatible_with = ["@platforms//os:linux"],
    visibility = ["@libpjrt_neuron//:__subpackages__"],
    deps = [":nki-cc"],
)

cc_library(
    name = "libpython",
    hdrs = ["libpython.h"],
    tags = ["manual"],
    visibility = ["//platforms/neuron:__subpackages__"],
    deps = [
        "@rules_python//python/cc:current_py_cc_headers",
        "@rules_python//python/cc:current_py_cc_libs",
    ],
)

cc_library(
    name = "empty",
    defines = ["ZML_RUNTIME_NEURON_DISABLED"],
    tags = ["manual"],
)

cc_library(
    name = "zmlxneuron",
    defines = ["ZML_RUNTIME_NEURON"],
    tags = ["manual"],
)

cc_library(
    name = "libnrt_headers",
    hdrs = ["nrt.h"],
    tags = ["manual"],
    visibility = [
        "//bin/zml-smi:__subpackages__",
    ],
    deps = ["@libpjrt_neuron//:libnrt_headers"],
)

cc_import(
    name = "neuron_ncfw",
    shared_library = "@aws-neuronx-runtime-lib//:libncfw.patchelf",
    tags = ["manual"],
)

cc_import(
    name = "neuron_nrt",
    shared_library = "@aws-neuronx-runtime-lib//:libnrt.patchelf",
    tags = ["manual"],
    deps = [":neuron_ncfw"],
)

zig_library(
    name = "python_launcher",
    import_name = "platforms/neuron/python_launcher",
    main = "python_launcher.zig",
    tags = ["manual"],
    visibility = ["//platforms/neuron:__subpackages__"],
    deps = [":libpython"],
)

zig_binary(
    name = "nki-cc",
    data = [
        "nki_kernel_compiler.py",
        "@neuron_py_deps//neuronx_cc",
        "@neuron_py_deps//nki",
    ],
    linkopts = ["-Wl,-rpath,$$ORIGIN/../lib"],
    main = "nki_compiler_launcher.zig",
    tags = ["manual"],
    target_compatible_with = ["@platforms//os:linux"],
    visibility = ["//platforms/neuron:__pkg__"],
    deps = [
        ":libpython",
        ":python_launcher",
    ],
)

zig_library(
    name = "nki_kernel",
    import_name = "platforms/neuron/nki_kernel",
    main = "nki_kernel.zig",
    tags = ["manual"],
    visibility = ["//zml:__pkg__"],
    deps = [
        ":neuron",
        "//bazel",
        "//stdx",
        "@rules_zig//zig/runfiles",
    ],
)

uv_lock(
    name = "requirements",
    srcs = ["requirements.in"],
    out = "requirements.lock.txt",
    args = [
        "--emit-index-url",
        "--emit-find-links",
        "--index-strategy=unsafe-best-match",
        "--upgrade",
        "--python-platform=x86_64-unknown-linux-gnu",
    ],
    tags = ["manual"],
)

zig_library(
    name = "neuron",
    import_name = "platforms/neuron",
    main = "neuron.zig",
    tags = ["manual"],
    visibility = ["//visibility:public"],
    deps = [
        "//pjrt",
    ] + select({
        "//platforms:neuron.enabled": [
            ":libnrt_headers",
            ":neuron_nrt",
            ":zmlxneuron",
            "//bazel",
            "//stdx",
            "@libpjrt_neuron",
            "@rules_zig//zig/runfiles",
        ],
        "//conditions:default": [":empty"],
    }),
)

zig_library(
    name = "topology",
    data = [":libpjrt_neuron"],
    import_name = "platforms/neuron/topology",
    main = "topology.zig",
    tags = ["manual"],
    visibility = ["//visibility:public"],
    deps = [
        ":neuron",
        ":zmlxneuron",
    ],
)