load("@bazel_skylib//rules:select_file.bzl", "select_file")
load("@rules_python//python:pip.bzl", "compile_pip_requirements")
load("@rules_python//python/entry_points:py_console_script_binary.bzl", "py_console_script_binary")
load("@rules_uv//uv:pip.bzl", "pip_compile")
load("@rules_zig//zig:defs.bzl", "zig_library")
load("@zml//bazel:cc_import.bzl", "cc_import")
load("@zml//bazel:runfiles.bzl", "runfiles_to_default")
load(":neuron.bzl", "py_binary_with_script")
load(":pyenv.bzl", "pyenv_zig")

pyenv_zig(
    name = "libneuronxla_pyenv_zig",
    deps = [":libneuronxla"],
)

zig_library(
    name = "libneuronxla_pyenv",
    main = ":libneuronxla_pyenv_zig",
)

pip_compile(
    name = "update_requirements",
    args = [
        "--generate-hashes",
        "--emit-index-url",
        "--emit-find-links",
        "--no-strip-extras",
        "--index-strategy unsafe-best-match",
        "--upgrade",
    ],
    python_platform = "x86_64-unknown-linux-gnu",
    requirements_in = "requirements.in",
    requirements_txt = "requirements.lock.txt",
)

py_console_script_binary(
    name = "neuronx-cc",
    binary_rule = py_binary_with_script,
    pkg = "@neuron_py_deps//neuronx_cc",
    visibility = ["//visibility:public"],
)

py_binary_with_script(
    name = "libneuronxla",
    srcs = ["empty.py"],
    main = "empty.py",
    deps = [
        ":neuronx-cc",
        "@neuron_py_deps//libneuronxla",
    ],
)

cc_library(
    name = "libpython",
    hdrs = ["libpython.h"],
    visibility = ["//visibility:public"],
    deps = [
        "@rules_python//python/cc:current_py_cc_headers",
        "@rules_python//python/cc:current_py_cc_libs",
    ],
)

compile_pip_requirements(
    name = "requirements",
    src = "requirements.in",
    py_binary = py_binary_with_script,
    requirements_txt = "requirements.lock.txt",
)

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

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

cc_import(
    name = "libpjrt_neuron_",
    add_needed = [
        "libpython3.11.so.1.0",
        "libzmlxneuron.so.0",
    ],
    data = [":libneuronxla"],
    shared_library = ":libneuronpjrt_so",
    soname = "libpjrt_neuron.so",
    deps = [
        ":zmlxneuron",
        "@aws-neuronx-runtime-lib",
        "@rules_python//python/cc:current_py_cc_libs",
    ],
)

cc_library(
    name = "zmlxneuron_lib",
    srcs = ["zmlxneuron.c"],
    linkopts = ["-ldl"],
)

cc_shared_library(
    name = "zmlxneuron_",
    shared_lib_name = "libzmlxneuron.so.0",
    deps = [":zmlxneuron_lib"],
)

cc_import(
    name = "zmlxneuron",
    shared_library = ":zmlxneuron_",
)

alias(
    name = "bootstrap",
    actual = "@rules_python//python/config_settings:bootstrap_impl",
)

cc_library(
    name = "empty",
)

cc_library(
    name = "libpjrt_neuron",
    defines = ["ZML_RUNTIME_NEURON"],
    deps = [":libpjrt_neuron_"],
)

zig_library(
    name = "neuron",
    import_name = "runtimes/neuron",
    main = "neuron.zig",
    visibility = ["//visibility:public"],
    deps = [
        "//pjrt",
    ] + select({
        "//runtimes:neuron.enabled": [
            ":libneuronxla_pyenv",
            ":libpjrt_neuron",
            ":libpython",
            "//async",
            "@rules_zig//zig/runfiles",
        ],
        "//conditions:default": [":empty"],
    }),
)

filegroup(
    name = "layers",
    srcs = [
        "@neuron_bookworm//bash/amd64",
        "@neuron_bookworm//libxml2/amd64",
        "@neuron_bookworm//zlib1g/amd64",
    ],
    visibility = ["//visibility:public"],
)