load("@bazel_lib//lib:copy_to_directory.bzl", "copy_to_directory")
load("@bazel_skylib//rules:select_file.bzl", "select_file")
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("@zml//bazel:patchelf.bzl", "patchelf")
load("@zml//bazel:runfiles.bzl", "runfiles_to_default")

runfiles_to_default(
    name = "libneuronxla_files",
    deps = ["@neuron_py_deps//libneuronxla:pkg"],
)

select_file(
    name = "libneuronpjrt_so_orig",
    srcs = ":libneuronxla_files",
    subpath = "site-packages/libneuronxla/libneuronpjrt.so",
)

patchelf(
    name = "libneuronpjrt_so",
    src = ":libneuronpjrt_so_orig",
    set_rpath = "$ORIGIN",
    soname = "libneuronpjrt.so",
)

patchelf(
    name = "libpjrt_neuron_so",
    src = "@zml//platforms/neuron:libpjrt_neuron",
    set_rpath = "$ORIGIN",
    soname = "libpjrt_neuron.so",
)

patchelf(
    name = "neuron_explorer",
    src = "@aws-neuronx-tools//:neuron-explorer",
    add_needed = ["libpjrt_neuron.so"],
    set_rpath = "$ORIGIN/../lib",
)

# This is a compiler-enabled sandbox. Python stays packaged because AWS exposes
# neuronx-cc and NKI compilation through Python entrypoints; ZML wraps those as
# native ELF launchers in bin/ so the image does not need shell scripts or host
# Python. Keep AWS compiler/runtime payloads intact and only trim Python
# files that are not part of Neuron compilation.
copy_to_directory(
    name = "sandbox",
    srcs = [
        ":libneuronpjrt_so",
        ":libneuronxla_files",
        ":libpjrt_neuron_so",
        ":neuron_explorer",
        "@aws-neuronx-collectives//:libnccom",
        "@aws-neuronx-runtime-lib//:libncfw.patchelf",
        "@aws-neuronx-runtime-lib//:libnrt.patchelf",
        "@aws-neuronx-runtime-lib//:libnrtucode_extisa",
        "@libgcc_s1",
        "@libgomp1",
        "@libstdcpp6",
        "@rules_python//python:current_py_toolchain",
        "@zlib1g",
        "@zml//platforms/neuron:neuronx-cc_files",
        "@zml//platforms/neuron:neuronx_cc_shim.py",
        "@zml//platforms/neuron:nki-cc_files",
        "@zml//platforms/neuron:nki_kernel_compiler.py",
    ],
    add_directory_to_runfiles = False,
    allow_overwrites = True,
    exclude_srcs_patterns = [
        "**/__pycache__/**",
        "**/*.pyc",
        "**/tests/**",
        "lib/libtcl*.so*",
        "lib/libtk*.so*",
        "lib/python*/idlelib/**",
        "lib/python*/tkinter/**",
        "lib/python*/turtledemo/**",
        "lib/python*/ensurepip/**",
        # "**/include/**",
        "**/_solib_*/**",
    ],
    hardlink = "on",
    include_external_repositories = ["**"],
    include_srcs_patterns = [
        "platforms/neuron/neuronx-cc",
        "platforms/neuron/neuronx_cc_shim.py",
        "platforms/neuron/nki-cc",
        "platforms/neuron/nki_kernel_compiler.py",
        "neuron_explorer/**",
        "**/lib**",
        "lib/python*/**",
        "site-packages/**",
    ],
    replace_prefixes = {
        "platforms/neuron/neuronx-cc": "bin/neuronx-cc",
        "platforms/neuron/neuronx_cc_shim.py": "bin/python-shims/neuronx-cc",
        "platforms/neuron/nki-cc": "bin/nki-cc",
        "platforms/neuron/nki_kernel_compiler.py": "site-packages/nki_kernel_compiler.py",
        "neuron_explorer": "bin",
        "libneuronpjrt_so": "lib",
        "libpjrt_neuron_so": "lib",
        "libnrt.patchelf": "lib",
        "libncfw.patchelf": "lib",
        "libnrtucode_extisa": "lib",
        "lib/x86_64-linux-gnu": "lib",
        "usr/lib/x86_64-linux-gnu": "lib",
        "opt/neuron": "lib",
    },
    visibility = [
        "@zml//tools/neuron:__pkg__",
        "@zml//zml:__pkg__",
        "@zml//zml/attention/nki:__pkg__",
    ],
)

cc_library(
    name = "libnrt_headers",
    linkopts = [
        # Defer function call resolution until the function is called
        # (lazy loading) rather than at load time.
        #
        # This is required because we want to let downstream use weak NRT symbols.
        #
        # We force it here because -z,now (which resolve all symbols at load time),
        # is the default in most bazel CC toolchains as well as in certain linkers.
        "-Wl,-z,lazy",
    ],
    visibility = ["@zml//platforms/neuron:__subpackages__"],
    deps = [
        "@aws-neuronx-runtime-lib//:libnrt_headers",
    ],
)

cc_library(
    name = "libpjrt_neuron",
    data = [":sandbox"],
    visibility = ["@zml//platforms/neuron:__subpackages__"],
)