load(
    "@com_github_johnynek_bazel_jar_jar//:jar_jar.bzl",
    "jar_jar",
)
load("@rules_java//java:java_binary.bzl", "java_binary")
load("@rules_java//java:java_import.bzl", "java_import")
load("@rules_java//java:java_library.bzl", "java_library")
load("@rules_java//java:java_test.bzl", "java_test")
load("@rules_pkg//pkg:mappings.bzl", "pkg_attributes", "pkg_files")
load("@rules_pkg//pkg:zip.bzl", "pkg_zip")
load("@rules_proto_grpc//java:defs.bzl", "java_proto_compile")
load("@rules_python//python:defs.bzl", "py_binary")
load("//bazel:ray.bzl", "define_java_module", "native_java_library")

exports_files([
    "testng.xml",
    "checkstyle.xml",
    "checkstyle-suppressions.xml",
])

all_modules = [
    "api",
    "runtime",
    "serve",
]

all_modules_with_test = all_modules + [
    "test",
    "performance_test",
]

java_import(
    name = "all_modules",
    jars = [
        "libio_ray_ray_" + module + ".jar"
        for module in all_modules
    ] + [
        "libio_ray_ray_" + module + "-src.jar"
        for module in all_modules
    ],
    deps = [
        ":io_ray_ray_" + module
        for module in all_modules
    ],
)

java_import(
    name = "all_modules_with_test",
    testonly = 1,
    jars = [
        "libio_ray_ray_" + module + ".jar"
        for module in all_modules_with_test
    ] + [
        "libio_ray_ray_" + module + "-src.jar"
        for module in all_modules_with_test
    ] + [
        "all_tests_bin_deploy.jar",
        "all_tests_bin_deploy-src.jar",
    ],
    deps = [
        ":io_ray_ray_" + module
        for module in all_modules_with_test
    ] + [
        ":all_tests",
    ],
)

define_java_module(
    name = "api",
    visibility = ["//visibility:public"],
    deps = [
        "@maven//:com_lmax_disruptor",
        "@maven//:com_sun_xml_bind_jaxb_core",
        "@maven//:com_sun_xml_bind_jaxb_impl",
        "@maven//:javax_activation_activation",
        "@maven//:javax_xml_bind_jaxb_api",
        "@maven//:org_apache_logging_log4j_log4j_api",
        "@maven//:org_apache_logging_log4j_log4j_core",
        "@maven//:org_apache_logging_log4j_log4j_slf4j_impl",
        "@maven//:org_slf4j_slf4j_api",
    ],
)

define_java_module(
    name = "runtime",
    additional_resources = [
        ":java_native_deps",
    ],
    additional_srcs = [
        ":all_java_proto",
    ],
    define_test_lib = True,
    exclude_srcs = [
        "runtime/src/main/java/io/ray/runtime/generated/*.java",
    ],
    test_deps = [
        ":io_ray_ray_api",
        ":io_ray_ray_runtime",
        "@maven//:commons_io_commons_io",
        "@maven//:javax_xml_bind_jaxb_api",
        "@maven//:org_apache_commons_commons_lang3",
        "@maven//:org_testng_testng",
    ],
    visibility = ["//visibility:public"],
    deps = [
        ":io_ray_ray_api",
        "@maven//:com_fasterxml_jackson_core_jackson_databind",
        "@maven//:com_github_java_json_tools_json_schema_validator",
        "@maven//:com_google_code_gson_gson",
        "@maven//:com_google_guava_guava",
        "@maven//:com_google_protobuf_protobuf_java",
        "@maven//:com_google_protobuf_protobuf_java_util",
        "@maven//:com_lmax_disruptor",
        "@maven//:com_typesafe_config",
        "@maven//:commons_io_commons_io",
        "@maven//:de_ruedigermoeller_fst",
        "@maven//:net_java_dev_jna_jna",
        "@maven//:org_apache_commons_commons_lang3",
        "@maven//:org_apache_logging_log4j_log4j_api",
        "@maven//:org_apache_logging_log4j_log4j_core",
        "@maven//:org_apache_logging_log4j_log4j_slf4j_impl",
        "@maven//:org_msgpack_msgpack_core",
        "@maven//:org_ow2_asm_asm",
        "@maven//:org_slf4j_slf4j_api",
        "@maven//:org_testng_testng",
    ],
)

define_java_module(
    name = "test",
    # (WangTaoTheTonic)For cpp x-lang tests. See //cpp:cluster_mode_xlang_test.
    visibility = ["//cpp:__subpackages__"],
    deps = [
        ":io_ray_ray_api",
        ":io_ray_ray_runtime",
        "@maven//:com_google_code_gson_gson",
        "@maven//:com_google_guava_guava",
        "@maven//:com_google_protobuf_protobuf_java",
        "@maven//:com_lmax_disruptor",
        "@maven//:com_sun_xml_bind_jaxb_core",
        "@maven//:com_sun_xml_bind_jaxb_impl",
        "@maven//:commons_io_commons_io",
        "@maven//:javax_xml_bind_jaxb_api",
        "@maven//:org_apache_commons_commons_lang3",
        "@maven//:org_apache_logging_log4j_log4j_api",
        "@maven//:org_apache_logging_log4j_log4j_core",
        "@maven//:org_apache_logging_log4j_log4j_slf4j_impl",
        "@maven//:org_slf4j_slf4j_api",
        "@maven//:org_testng_testng",
    ],
)

define_java_module(
    name = "performance_test",
    deps = [
        ":io_ray_ray_api",
        ":io_ray_ray_runtime",
        "@maven//:com_google_code_gson_gson",
        "@maven//:com_google_guava_guava",
        "@maven//:com_lmax_disruptor",
        "@maven//:commons_io_commons_io",
        "@maven//:org_apache_commons_commons_lang3",
        "@maven//:org_apache_logging_log4j_log4j_api",
        "@maven//:org_apache_logging_log4j_log4j_core",
        "@maven//:org_apache_logging_log4j_log4j_slf4j_impl",
        "@maven//:org_slf4j_slf4j_api",
    ],
)

define_java_module(
    name = "serve",
    additional_srcs = [
        ":serve_java_proto",
    ],
    define_test_lib = True,
    exclude_srcs = [
        "serve/src/main/java/io/ray/serve/generated/*.java",
    ],
    test_deps = [
        ":io_ray_ray_api",
        ":io_ray_ray_runtime",
        ":io_ray_ray_serve",
        "@maven//:com_google_code_gson_gson",
        "@maven//:com_google_guava_guava",
        "@maven//:com_google_protobuf_protobuf_java",
        "@maven//:commons_io_commons_io",
        "@maven//:org_apache_commons_commons_lang3",
        "@maven//:org_apache_httpcomponents_client5_httpclient5",
        "@maven//:org_apache_httpcomponents_client5_httpclient5_fluent",
        "@maven//:org_apache_httpcomponents_core5_httpcore5",
        "@maven//:org_slf4j_slf4j_api",
        "@maven//:org_testng_testng",
    ],
    visibility = ["//visibility:public"],
    deps = [
        ":io_ray_ray_api",
        ":io_ray_ray_runtime",
        "@maven//:com_google_code_gson_gson",
        "@maven//:com_google_guava_guava",
        "@maven//:com_google_protobuf_protobuf_java",
        "@maven//:org_apache_commons_commons_lang3",
        "@maven//:org_apache_httpcomponents_core5_httpcore5",
        "@maven//:org_slf4j_slf4j_api",
    ],
)

java_library(
    name = "all_tests_lib",
    runtime_deps = [
        ":io_ray_ray_performance_test",
        ":io_ray_ray_runtime_test",
        ":io_ray_ray_serve_test",
        ":io_ray_ray_test",
    ],
)

# This is a local java test rule. It needs generated files to be copied into
# the source tree before running. To build and generate ray core (gcs and
# raylet), run `bazelisk run //:gen_ray_pkg` first.
#
# This rule used to depend on local genrules, which are deprecated. Reason
# being that local genrules are build rules and do not capture changes in
# source tree, and hence cannot be cached by bazel remote cache. Using local
# genrule forces bazel to effectively disable caching globally to have a
# correct build.
#
# TODO(ray-ci): covert java tests to non-local, hermetic tests.
java_test(
    name = "all_tests",
    testonly = True,
    args = ["java/testng.xml"],
    data = [
        "testng.xml",
        "//:ray_pkg_zip",
    ],
    main_class = "org.testng.TestNG",
    resources = [
        "//cpp:counter.so",
        "//cpp:plus.so",
    ],
    tags = ["local"],
    runtime_deps = [
        ":all_tests_lib",
    ],
)

# java_test does not produce a _deploy.jar in Bazel 7+. Use a java_binary
# companion so that all_tests_bin_deploy.jar is available in both Bazel 6 and 7.
java_binary(
    name = "all_tests_bin",
    testonly = True,
    main_class = "org.testng.TestNG",
    resources = [
        "//cpp:counter.so",
        "//cpp:plus.so",
    ],
    runtime_deps = [":all_tests_lib"],
)

# 0. `cp testng_custom_template.xml testng_custom.xml`
# 1. `bazel run //:gen_ray_pkg`
# 2. Specify test class/method in `testng_custom.xml`
# 3. `bazel test //java:custom_test --test_output=streamed`
java_test(
    name = "custom_test",
    args = ["java/testng_custom.xml"],
    data = [
        "testng_custom.xml",
        "//:ray_pkg_zip",
    ],
    main_class = "org.testng.TestNG",
    tags = ["local"],
    runtime_deps = [
        ":all_tests_lib",
    ],
)

# We'd better make resource files can be accessed from 3rd party library.
# More detail please see https://github.com/ray-project/ray/pull/21641.
java_proto_compile(
    name = "common_java_proto",
    deps = ["@io_ray//src/ray/protobuf:common_proto"],
)

java_proto_compile(
    name = "runtime_env_common_java_proto",
    deps = [
        "@io_ray//src/ray/protobuf:runtime_env_common_proto",
        "@io_ray//src/ray/protobuf/public:runtime_environment_proto",
    ],
)

java_proto_compile(
    name = "gcs_java_proto",
    deps = ["@io_ray//src/ray/protobuf:gcs_proto"],
)

java_proto_compile(
    name = "serve_java_proto",
    deps = ["@io_ray//src/ray/protobuf:serve_proto"],
)

filegroup(
    name = "all_java_proto",
    srcs = [
        ":common_java_proto",
        ":gcs_java_proto",
        ":runtime_env_common_java_proto",
    ],
)

native_java_library(
    name = "core_worker_library_java",
    module_name = "runtime",
    native_library_name = "//src/ray/core_worker/lib/java:libcore_worker_library_java.so",
)

filegroup(
    name = "java_native_deps",
    srcs = [
        ":core_worker_library_java",
    ],
)

pkg_files(
    name = "api_pom_files",
    srcs = ["io_ray_ray_api_pom"],
    prefix = "api/",
    renames = {
        "io_ray_ray_api_pom.xml": "pom.xml",
    },
    visibility = ["//visibility:private"],
)

pkg_files(
    name = "runtime_pom_files",
    srcs = ["io_ray_ray_runtime_pom"],
    prefix = "runtime/",
    renames = {
        "io_ray_ray_runtime_pom.xml": "pom.xml",
    },
    visibility = ["//visibility:private"],
)

pkg_files(
    name = "test_pom_files",
    srcs = ["io_ray_ray_test_pom"],
    prefix = "test/",
    renames = {
        "io_ray_ray_test_pom.xml": "pom.xml",
    },
    visibility = ["//visibility:private"],
)

pkg_files(
    name = "performance_test_pom_files",
    srcs = ["io_ray_ray_performance_test_pom"],
    prefix = "performance_test/",
    renames = {
        "io_ray_ray_performance_test_pom.xml": "pom.xml",
    },
    visibility = ["//visibility:private"],
)

pkg_files(
    name = "serve_pom_files",
    srcs = ["io_ray_ray_serve_pom"],
    prefix = "serve/",
    renames = {
        "io_ray_ray_serve_pom.xml": "pom.xml",
    },
    visibility = ["//visibility:private"],
)

pkg_zip(
    name = "pom_files",
    srcs = [
        ":api_pom_files",
        ":performance_test_pom_files",
        ":runtime_pom_files",
        ":serve_pom_files",
        ":test_pom_files",
    ],
    visibility = ["//visibility:private"],
)

py_binary(
    name = "gen_pom_files",
    srcs = ["gen_pom_files.py"],
    data = [":pom_files.zip"],
    visibility = ["//visibility:private"],
    deps = ["//bazel:gen_extract"],
)

# Generates the dependencies needed by maven.
genrule(
    name = "proto_files",
    srcs = [
        ":all_java_proto",
        ":serve_java_proto",
    ],
    outs = ["proto_files.zip"],
    cmd = """
        set -euo pipefail

        tmpdir=$$(mktemp -d)

        mkdir -p "$$tmpdir/java/runtime/src/main/java/io/ray/runtime/generated"
        for f in $(locations :all_java_proto); do
            unzip -q "$$f" -x META-INF/MANIFEST.MF -d "$$tmpdir/java/runtime/src/main/java"
        done

        mkdir -p "$$tmpdir/java/serve/src/main/java/io/ray/serve/generated"
        for f in $(locations :serve_java_proto); do
            unzip -q "$$f" -x META-INF/MANIFEST.MF -d "$$tmpdir/java/serve/src/main/java"
        done

        (cd "$$tmpdir/java"; zip -0 -q -r out.zip runtime serve)
        mv "$$tmpdir/java/out.zip" $@

        rm -rf "$$tmpdir"
    """,
    visibility = ["//visibility:private"],
)

py_binary(
    name = "gen_proto_files",
    srcs = ["gen_proto_files.py"],
    data = [":proto_files.zip"],
    visibility = ["//visibility:private"],
    deps = ["//bazel:gen_extract"],
)

pkg_files(
    name = "maven_deps_files",
    srcs = [
        ":java_native_deps",
    ],
    attributes = pkg_attributes(mode = "755"),
    prefix = select(
        {
            "@platforms//os:linux": "runtime/native_dependencies/native/linux",
            "@platforms//os:macos": "runtime/native_dependencies/native/darwin",
        },
        no_match_error = "Unsupported platform",
    ),
    visibility = ["//visibility:private"],
)

pkg_zip(
    name = "maven_deps",
    srcs = [
        ":maven_deps_files",
    ],
    visibility = ["//visibility:private"],
)

py_binary(
    name = "gen_maven_deps",
    srcs = ["gen_maven_deps.py"],
    data = [":maven_deps.zip"],
    visibility = ["//visibility:private"],
    deps = ["//bazel:gen_extract"],
)

java_binary(
    name = "ray_dist",
    # This rule is used to package all Ray Java code and the third-party dependencies into a
    # fat jar file. It's not really an executable jar. So we set its `main_class` to empty.
    main_class = "",
    runtime_deps = [
        "//java:io_ray_ray_api",
        "//java:io_ray_ray_runtime",
        "//java:io_ray_ray_serve",
    ],
)

jar_jar(
    name = "ray_dist_shaded",
    input_jar = "//java:ray_dist_deploy.jar",
    rules = "//java:shade_rule",
)

# Shade dependencies in tests fat jar.
jar_jar(
    name = "all_tests_shaded",
    input_jar = "//java:all_tests_bin_deploy.jar",
    rules = "//java:shade_rule",
)

pkg_files(
    name = "ray_java_pkg_files",
    srcs = [
        ":ray_dist_shaded.jar",
    ],
    prefix = "ray/jars/",
    renames = {
        "ray_dist_shaded.jar": "ray_dist.jar",
    },
)

pkg_zip(
    name = "ray_java_pkg_zip",
    srcs = [
        ":ray_java_pkg_files",
    ],
    out = "ray_java_pkg.zip",
    visibility = ["//visibility:private"],
)

py_binary(
    name = "gen_ray_java_pkg",
    srcs = ["gen_ray_java_pkg.py"],
    data = [
        ":ray_java_pkg.zip",
    ],
    visibility = ["//visibility:private"],
    deps = [
        "//bazel:gen_extract",
    ],
)