"""Generate zipped aar file including different variants of .so in jni folder."""

load("@build_bazel_rules_android//android:rules.bzl", "android_binary")

def aar_with_jni(
        name,
        android_library,
        headers = None,
        flatten_headers = False,
        strip_headers_prefix = "",
        license_file = "//:LICENSE",
        third_party_notice = None):
    """Generates an Android AAR with repo root license given an Android library target.

    Args:
      name: Name of the generated .aar file.
      android_library: The `android_library` target to package. Note that the
          AAR will contain *only that library's .jar` sources. It does not
          package the transitive closure of all Java source dependencies.
      headers: Optional list of headers that will be included in the
          generated .aar file. This is useful for distributing self-contained
          .aars with native libs that can be used directly by native clients.
      flatten_headers: Whether to flatten the output paths of included headers.
      strip_headers_prefix: The prefix to strip from the output paths of included headers.
      license_file: Optional. The main LICENSE file to include in the AAR.
          Defaults to //third_party/tensorflow:LICENSE.
      third_party_notice: Optional. The third party dependency licenses as THIRD_PARTY_NOTICE.txt.
    """

    # Generate dummy AndroidManifest.xml for dummy apk usage
    # (dummy apk is generated by <name>_dummy_app_for_so target below)
    native.genrule(
        name = name + "_binary_manifest_generator",
        outs = [name + "_generated_AndroidManifest.xml"],
        cmd = """
cat > $(OUTS) <<EOF
<manifest
  xmlns:android="http://schemas.android.com/apk/res/android"
  package="dummy.package.for.so">
  <uses-sdk android:minSdkVersion="34"/>
</manifest>
EOF
""",
    )

    # Generate dummy apk including .so files and later we extract out
    # .so files and throw away the apk.
    android_binary(
        name = name + "_dummy_app_for_so",
        manifest = name + "_generated_AndroidManifest.xml",
        custom_package = "dummy.package.for.so",
        deps = [android_library],
        # In some platforms we don't have an Android SDK/NDK and this target
        # can't be built. We need to prevent the build system from trying to
        # use the target in that case.
        tags = [
            "manual",
            "no_cuda_on_cpu_tap",
        ],
    )

    srcs = [
        android_library + ".aar",
        name + "_dummy_app_for_so_unsigned.apk",
        license_file,
    ]

    cmd = """
cp $(location {0}.aar) $(location :{1}.aar)
chmod +w $(location :{1}.aar)
origdir=$$PWD
cd $$(mktemp -d)
unzip $$origdir/$(location :{1}_dummy_app_for_so_unsigned.apk) "lib/*"
cp -r lib jni
zip -r $$origdir/$(location :{1}.aar) jni/*/*.so
cp $$origdir/$(location {2}) ./LICENSE
zip $$origdir/$(location :{1}.aar) LICENSE
""".format(android_library, name, license_file)

    if headers:
        srcs += headers
        cmd += """
        mkdir headers
        """
        for src in headers:
            if flatten_headers:
                cmd += """
                    cp -RL $$origdir/$(location {0}) headers/$$(basename $(location {0}))
                """.format(src)
            else:
                cmd += """
                    default_dir=$$(dirname $(rootpath {0}))
                    modified_dir=$$(echo $$default_dir | sed -e 's/^{1}//g')
                    mkdir -p headers/$$modified_dir
                    cp -RL $$origdir/$(location {0}) headers/$$modified_dir
                    if [ -n "{1}" ]; then
                      sed -i -e 's/^#include \"{1}/#include \"/g' headers/$$modified_dir/$$(basename $(location {0}))
                    fi
                """.format(src, strip_headers_prefix.replace("/", "\\/"))
        cmd += "zip -r $$origdir/$(location :{0}.aar) headers".format(name)

    if third_party_notice:
        srcs.append(third_party_notice)
        cmd += """
            cp $$origdir/$(location {0}) ./THIRD_PARTY_NOTICE.txt
            zip $$origdir/$(location :{1}.aar) THIRD_PARTY_NOTICE.txt
        """.format(third_party_notice, name)

    native.genrule(
        name = name,
        srcs = srcs,
        outs = [name + ".aar"],
        # In some platforms we don't have an Android SDK/NDK and this target
        # can't be built. We need to prevent the build system from trying to
        # use the target in that case.
        tags = ["manual"],
        cmd = cmd,
    )

def aar_without_jni(
        name,
        android_library,
        license_file = "//:LICENSE"):
    """Generates an Android AAR with repo root license given a pure Java Android library target.

    Args:
      name: Name of the generated .aar file.
      android_library: The `android_library` target to package. Note that the
          AAR will contain *only that library's .jar` sources. It does not
          package the transitive closure of all Java source dependencies.
      license_file: Optional. The main LICENSE file to include in the AAR.
          Defaults to //third_party/tensorflow:LICENSE.
    """

    srcs = [
        android_library + ".aar",
        license_file,
    ]

    cmd = """
cp $(location {0}.aar) $(location :{1}.aar)
chmod +w $(location :{1}.aar)
origdir=$$PWD
cd $$(mktemp -d)
cp $$origdir/$(location {2}) ./LICENSE
zip $$origdir/$(location :{1}.aar) LICENSE
""".format(android_library, name, license_file)

    native.genrule(
        name = name,
        srcs = srcs,
        outs = [name + ".aar"],
        cmd = cmd,
    )