# Copyright 2022 The KerasCV Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
# Usage: python configure.py
"""Configures local environment to prepare for building KerasCV from source."""


import logging
import os
import pathlib
import platform

import tensorflow as tf
from packaging.version import Version

_TFA_BAZELRC = ".bazelrc"


# Writes variables to bazelrc file
def write(line):
    with open(_TFA_BAZELRC, "a") as f:
        f.write(line + "\n")


def write_action_env(var_name, var):
    write('build --action_env {}="{}"'.format(var_name, var))


def is_macos():
    return platform.system() == "Darwin"


def is_windows():
    return platform.system() == "Windows"


def is_linux():
    return platform.system() == "Linux"


def is_raspi_arm():
    return os.uname()[4] == "armv7l" or os.uname()[4] == "aarch64"


def is_linux_ppc64le():
    return is_linux() and platform.machine() == "ppc64le"


def is_linux_x86_64():
    return is_linux() and platform.machine() == "x86_64"


def is_linux_arm():
    return is_linux() and platform.machine() == "arm"


def is_linux_aarch64():
    return is_linux() and platform.machine() == "aarch64"


def is_linux_s390x():
    return is_linux() and platform.machine() == "s390x"


def get_tf_header_dir():
    import tensorflow as tf

    tf_header_dir = tf.sysconfig.get_compile_flags()[0][2:]
    if is_windows():
        tf_header_dir = tf_header_dir.replace("\\", "/")
    return tf_header_dir


def get_cpp_version():
    cpp_version = "c++14"
    if Version(tf.__version__) >= Version("2.10"):
        cpp_version = "c++17"
    return cpp_version


def get_tf_shared_lib_dir():
    import tensorflow as tf

    # OS Specific parsing
    if is_windows():
        tf_shared_lib_dir = tf.sysconfig.get_compile_flags()[0][2:-7] + "python"
        return tf_shared_lib_dir.replace("\\", "/")
    elif is_raspi_arm():
        return tf.sysconfig.get_compile_flags()[0][2:-7] + "python"
    else:
        return tf.sysconfig.get_link_flags()[0][2:]


# Converts the linkflag namespec to the full shared library name
def get_shared_lib_name():
    import tensorflow as tf

    namespec = tf.sysconfig.get_link_flags()
    if is_macos():
        # MacOS
        return "lib" + namespec[1][2:] + ".dylib"
    elif is_windows():
        # Windows
        return "_pywrap_tensorflow_internal.lib"
    elif is_raspi_arm():
        # The below command for linux would return an empty list
        return "_pywrap_tensorflow_internal.so"
    else:
        # Linux
        return namespec[1][3:]


def create_build_configuration():
    print()
    print("Configuring KerasCV to be built from source...")

    if os.path.isfile(_TFA_BAZELRC):
        os.remove(_TFA_BAZELRC)

    logging.disable(logging.WARNING)

    write_action_env("TF_HEADER_DIR", get_tf_header_dir())
    write_action_env("TF_SHARED_LIBRARY_DIR", get_tf_shared_lib_dir())
    write_action_env("TF_SHARED_LIBRARY_NAME", get_shared_lib_name())
    write_action_env("TF_CXX11_ABI_FLAG", tf.sysconfig.CXX11_ABI_FLAG)

    # This should be replaced with a call to tf.sysconfig if it's added
    write_action_env("TF_CPLUSPLUS_VER", get_cpp_version())

    write("build --spawn_strategy=standalone")
    write("build --strategy=Genrule=standalone")
    write("build  --experimental_repo_remote_exec")
    write("build -c opt")
    write(
        "build --cxxopt="
        + '"-D_GLIBCXX_USE_CXX11_ABI="'
        + str(tf.sysconfig.CXX11_ABI_FLAG)
    )

    if is_windows():
        write("build --config=windows")
        write("build:windows --enable_runfiles")
        write("build:windows --copt=/experimental:preprocessor")
        write("build:windows --host_copt=/experimental:preprocessor")
        write("build:windows --copt=/arch=AVX")
        write("build:windows --cxxopt=/std:" + get_cpp_version())
        write("build:windows --host_cxxopt=/std:" + get_cpp_version())

    if is_macos() or is_linux():
        if not is_linux_ppc64le() and not is_linux_arm() and not is_linux_aarch64():
            write("build --copt=-mavx")
        write("build --cxxopt=-std=" + get_cpp_version())
        write("build --host_cxxopt=-std=" + get_cpp_version())

    print("> Building only CPU ops")

    print()
    print("Build configurations successfully written to", _TFA_BAZELRC, ":\n")
    print(pathlib.Path(_TFA_BAZELRC).read_text())


if __name__ == "__main__":
    create_build_configuration()