# Copyright 2021 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

# Build/install the iree-compiler-backend python package.
# Note that this includes a relatively large build of LLVM (~2400 C++ files)
# and can take a considerable amount of time, especially with defaults.
# To install:
#   pip install . --use-feature=in-tree-build
# To build a wheel:
#   pip wheel . --use-feature=in-tree-build
#
# It is recommended to build with Ninja and ccache. To do so, set environment
# variables by prefixing to above invocations:
#   CMAKE_C_COMPILER_LAUNCHER=ccache CMAKE_CXX_COMPILER_LAUNCHER=ccache
#
# On CIs, it is often advantageous to re-use/control the CMake build directory.
# This can be set with the IREE_COMPILER_API_CMAKE_BUILD_DIR env var.
import json
import os
import platform
import re
import shutil
import subprocess
import sys
import sysconfig

from distutils.command.build import build as _build
from setuptools import find_namespace_packages, setup, Extension
from setuptools.command.build_ext import build_ext as _build_ext
from setuptools.command.build_py import build_py as _build_py

# Setup and get version information.
THIS_DIR = os.path.realpath(os.path.dirname(__file__))
IREESRC_DIR = os.path.join(THIS_DIR, "..", "..")
VERSION_INFO_FILE = os.path.join(IREESRC_DIR, "version_info.json")


def load_version_info():
  with open(VERSION_INFO_FILE, "rt") as f:
    return json.load(f)


try:
  version_info = load_version_info()
except FileNotFoundError:
  print("version_info.json not found. Using defaults")
  version_info = {}

PACKAGE_SUFFIX = version_info.get("package-suffix") or ""
PACKAGE_VERSION = version_info.get("package-version") or "0.1dev1"


class CustomBuild(_build):

  def run(self):
    self.run_command("build_py")
    self.run_command("build_ext")
    self.run_command("build_scripts")


class CMakeExtension(Extension):

  def __init__(self, name, sourcedir=""):
    Extension.__init__(self, name, sources=[])
    self.sourcedir = os.path.abspath(sourcedir)


class CMakeBuildPy(_build_py):

  def run(self):
    version_py_content = generate_version_py()
    print(f"Generating version.py:\n{version_py_content}", file=sys.stderr)
    subprocess.check_call(["cmake", "--version"])

    target_dir = os.path.abspath(self.build_lib)
    print(f"Building in target dir: {target_dir}", file=sys.stderr)
    os.makedirs(target_dir, exist_ok=True)
    cmake_build_dir = os.getenv("IREE_COMPILER_API_CMAKE_BUILD_DIR")
    if not cmake_build_dir:
      cmake_build_dir = os.path.join(target_dir, "..", "cmake_build")
    os.makedirs(cmake_build_dir, exist_ok=True)
    cmake_build_dir = os.path.abspath(cmake_build_dir)
    print(f"CMake build dir: {cmake_build_dir}", file=sys.stderr)
    cmake_install_dir = os.path.abspath(
        os.path.join(target_dir, "..", "cmake_install"))
    print(f"CMake install dir: {cmake_install_dir}", file=sys.stderr)
    src_dir = os.path.abspath(os.path.dirname(__file__))
    cfg = "Release"
    cmake_args = [
        "-GNinja",
        "--log-level=VERBOSE",
        "-DCMAKE_INSTALL_PREFIX={}".format(cmake_install_dir),
        "-DPython3_EXECUTABLE={}".format(sys.executable),
        "-DPython3_INCLUDE_DIRS={}".format(sysconfig.get_path("include")),
        "-DIREE_VERSION_INFO={}".format(self.distribution.get_version()),
        "-DCMAKE_BUILD_TYPE={}".format(cfg),
    ]

    # Enable CUDA if specified.
    cuda_target_option = os.getenv("IREE_TARGET_BACKEND_CUDA")
    if cuda_target_option:
      cmake_args.append(f"-DIREE_TARGET_BACKEND_CUDA={cuda_target_option}")

    build_args = []
    if os.path.exists(cmake_install_dir):
      shutil.rmtree(cmake_install_dir)
    cmake_cache_file = os.path.join(cmake_build_dir, "CMakeCache.txt")
    if os.path.exists(cmake_cache_file):
      os.remove(cmake_cache_file)
    install_target = "install/strip"
    if platform.system() == "Windows":
      install_target = "install"
    print(f"Configuring with: {cmake_args}", file=sys.stderr)
    subprocess.check_call(["cmake", src_dir] + cmake_args, cwd=cmake_build_dir)
    subprocess.check_call(
        ["cmake", "--build", ".", "--target", install_target] + build_args,
        cwd=cmake_build_dir)
    print("Build complete.", file=sys.stderr)
    if os.path.exists(target_dir):
      shutil.rmtree(target_dir)
    print("Copying install to target.", file=sys.stderr)
    shutil.copytree(os.path.join(cmake_install_dir, "python_package"),
                    target_dir,
                    symlinks=False)
    print("Target populated.", file=sys.stderr)

    # Write version.py.
    version_py_file = os.path.join(target_dir, "iree", "compiler", "version.py")
    os.makedirs(os.path.dirname(version_py_file), exist_ok=True)
    with open(version_py_file, "wt") as f:
      f.write(version_py_content)


class NoopBuildExtension(_build_ext):

  def __init__(self, *args, **kwargs):
    assert False

  def build_extension(self, ext):
    pass


def generate_version_py():
  return f"""# Auto-generated version info.
PACKAGE_SUFFIX = "{PACKAGE_SUFFIX}"
VERSION = "{PACKAGE_VERSION}"
REVISIONS = {json.dumps(find_git_versions())}
"""


def find_git_versions():
  revisions = {}
  try:
    revisions["IREE"] = subprocess.check_output(
        ["git", "rev-parse", "HEAD"], cwd=IREESRC_DIR).decode("utf-8").strip()
  except subprocess.SubprocessError as e:
    print(f"ERROR: Could not get IREE revision: {e}", file=sys.stderr)
  revisions["LLVM_PROJECT"] = find_git_submodule_revision(
      "third_party/llvm-project")
  revisions["TENSORFLOW"] = find_git_submodule_revision(
      "third_party/tensorflow")
  revisions["MLIR_HLO"] = find_git_submodule_revision("third_party/mlir-hlo")
  return revisions


def find_git_submodule_revision(submodule_path):
  try:
    data = subprocess.check_output(["git", "ls-tree", "HEAD", submodule_path],
                                   cwd=IREESRC_DIR).decode("utf-8").strip()
    columns = re.split("\\s+", data)
    return columns[2]
  except Exception as e:
    print(
        f"ERROR: Could not get submodule revision for {submodule_path}"
        f" ({e})",
        file=sys.stderr)
    return ""


setup(
    name=f"iree-compiler{PACKAGE_SUFFIX}",
    version=f"{PACKAGE_VERSION}",
    author="IREE Authors",
    author_email="iree-discuss@googlegroups.com",
    description="IREE Compiler API",
    long_description="",
    license="Apache-2.0",
    classifiers=[
        "Programming Language :: Python :: 3.7",
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
    ],
    ext_modules=[
        CMakeExtension("iree.compiler._mlir_libs._mlir"),
        CMakeExtension("iree.compiler._mlir_libs._ireeDialects"),
        CMakeExtension("iree.compiler._mlir_libs._ireecTransforms"),
        CMakeExtension("iree.compiler._mlir_libs._mlirHlo"),
        CMakeExtension("iree.compiler._mlir_libs._mlirLinalgPasses"),
    ],
    cmdclass={
        "build": CustomBuild,
        "built_ext": NoopBuildExtension,
        "build_py": CMakeBuildPy,
    },
    zip_safe=False,
    packages=find_namespace_packages(include=[
        "iree.compiler",
        "iree.compiler.*",
    ],),
    entry_points={
        "console_scripts": [
            "ireec = iree.compiler.tools.scripts.ireec.__main__:main",
            # Transitional note: iree-translate resolves to ireec.
            "iree-translate = iree.compiler.tools.scripts.ireec.__main__:main",
        ],
    },
    install_requires=[
        "numpy",
        "PyYAML",
    ],
)