"""
Final build of BornAgain, with multiple versions of Python.

To be run manually on the packaging machine.

Results need to be checked manually, as things may go wrong silently.
"""

import os, platform, re, subprocess as subp
import re
from timeit import default_timer as timer
from collections import namedtuple

# all environmental variables
ENV = os.environ
BUILD_DIR = "build"

# separator line
SEPLINE = "-" * 40

# directories needed for the full build
BuildDirs = namedtuple("BuildDirectories", ["source", "build", "package", "python"])

class WindowsMake:
    """ Build on Windows """
    PLATFORM = "windows"
    ARCH = "amd64"
    NPROC = 8
    # initial path
    PATH_INI = ENV["PATH"]
    OPT_DIR = "C:/opt/x64"
    FFTW3_INCLUDE_DIR = OPT_DIR + "/include"
    FFTW3_LIB = OPT_DIR + "/lib/libfftw3-3.lib"
    BOOST_DIR = OPT_DIR + "/boost_current"
    BOOST_INCLUDE_DIR = BOOST_DIR + "/include"
    BOOST_LIB_DIR = BOOST_DIR + "/lib"
    QT_MSVC_DIR = "C:/Qt/6.2.4/msvc2019_64"
    QTCMake_DIR = QT_MSVC_DIR + "/lib/cmake"
    PY_PLATFORM_BASE = "C:/opt/multipython/Python"
    # supported Python platforms: default (3.9), 3.8, 3.10, 3.11
    SUPPORTED_PY_VERSIONS = ("", "38", "310", "311")

    def configure(dirs:BuildDirs):
        """
        Configure the build with a given source and build directory
        and a given Python platform
        """

        # NOTE: The command is equivalent to the following shell command:
        # $ cmake -G "Visual Studio 17 2022" -A x64 -T host=x64 ... \
        #     -DCMAKE_PREFIX_PATH="optdir;qtcmakedir" ...
        # Notice that the double-quotation marks should not be included
        # in the passed arguments.
        conf_prc = subp.run(
            ["cmake", "-G", "Visual Studio 17 2022",
             "-A x64",
             "-T host=x64",
             "-DCMAKE_PREFIX_PATH=" + WindowsMake.OPT_DIR + ";" + WindowsMake.QTCMake_DIR,
             "-DFFTW3_INCLUDE_DIR=" + WindowsMake.FFTW3_INCLUDE_DIR,
             "-DFFTW3_LIBRARY=" + WindowsMake.FFTW3_LIB,
             "-DCMAKE_INCLUDE_PATH=" + WindowsMake.OPT_DIR + "/include;"
             + WindowsMake.BOOST_INCLUDE_DIR,
             "-DCMAKE_LIBRARY_PATH=" + WindowsMake.OPT_DIR + "/lib;"
             + WindowsMake.BOOST_LIB_DIR,
             "-DBA_PY_PLATFORM=" + dirs.python,
             "-DCMAKE_C_COMPILER=cl.exe",
             "-DCMAKE_CXX_COMPILER=cl.exe",
             "-B", dirs.build, "-S", dirs.source])

        return conf_prc

    def build(dirs:BuildDirs):
        """ Build in the given build directory """
        build_prc = subp.run(["cmake", "--build", dirs.build, "--config Release", "--",
                              "/fl", "/flp:logfile=BornAgainMSBuild.log", "/verbosity:minimal"])
        return build_prc

    def test(dirs:BuildDirs):
        """ Perform tests """
        os.chdir(dirs.build)
        # Windows: change the system PATH temporarily
        _PATH = WindowsMake.QT_MSVC_DIR + "/bin;"  + dirs.python + ";" + WindowsMake.PATH_INI
        test_prc = subp.run(["ctest", "-C", "Release",
                             "--parallel %i" % WindowsMake.NPROC, "--output-on-failure"],
                             env={"PATH": _PATH})
        return test_prc

    def pack(dirs:BuildDirs):
        """ Build package """
        pack_prc = subp.run(["cpack", dirs.build, "-C", "Release", "-B", dirs.package])
        return pack_prc


class LinuxMake:
    PLATFORM = "linux"
    ARCH = "x86_64"
    CC = "gcc"
    CXX = "g++"
    MPLBACKEND = "Agg"
    QTCMake_DIR = "/usr/local/Qt6/6.2.3/gcc_64/lib/cmake"
    NPROC = 8
    PY_PLATFORM_BASE = "/home/build/multipython/Python"
    # supported Python platforms: 3.9, 3.8, 3.10, 3.11
    SUPPORTED_PY_VERSIONS = ("39", "38", "310", "311")

    def _setenv():
        ENV["CC"] = LinuxMake.CC
        ENV["CXX"] = LinuxMake.CXX
        ENV["MPLBACKEND"] = LinuxMake.MPLBACKEND

    def configure(dirs:BuildDirs):
        """
        Configure the build with a given source and build directory
        and a given Python platform
        """
        LinuxMake._setenv()
        conf_prc = subp.run(
            ["cmake",
             "-DCMAKE_PREFIX_PATH=" + LinuxMake.QTCMake_DIR,
             "-DZERO_TOLERANCE=ON",
             "-DDEVELOPER_CHECKS=ON",
             "-DBA_PY_PLATFORM=" + dirs.python,
             "-B", dirs.build, "-S", dirs.source])

        return conf_prc

    def build(dirs:BuildDirs):
        """ Build in the given build directory """
        build_prc = subp.run(["cmake", "--build", dirs.build, "--config Release",
                              "--parallel %i" % LinuxMake.NPROC])
        return build_prc

    def test(dirs:BuildDirs):
        """ Perform tests """
        os.chdir(dirs.build)
        test_prc = subp.run(["ctest", "-C", "Release", "--parallel %i" % LinuxMake.NPROC,
                             "--output-on-failure"])
        return test_prc

    def pack(dirs:BuildDirs):
        """ Build package """
        wheel_prc = subp.run(["bash", dirs.build + "/var/mk_wheel_multilinux.sh"])
        wheel_prc.check_returncode()
        pack_prc = subp.run(["cpack", dirs.build, "-B " + dirs.package])
        return pack_prc


class MacMake_x64:
    PLATFORM = "darwin"
    ARCH = "x86_64"
    QT_DIR = "/usr/local/opt/qt"
    QTCMake_DIR = QT_DIR + "/lib/cmake"
    OPTDIR="/Users/Shared/Software"
    # nr of logical CPUs can be obtained via `sysctl hw.logicalcpu`
    NPROC = 6
    PY_PLATFORM_BASE = "/opt/multipython/Python"
    # supported Python platforms: default (3.9), 3.8, 3.10, 3.11
    SUPPORTED_PY_VERSIONS = ("", "38", "310", "311")

    def configure(dirs:BuildDirs):
        """
        Configure the build with a given source and build directory
        and a given Python platform
        """

        conf_prc = subp.run(
            ["cmake",
             "-DCMAKE_PREFIX_PATH=" + MacMake_x64.OPTDIR + ";" + MacMake_x64.QTCMake_DIR,
             "-DCMAKE_BUILD_TYPE=Release",
             "-DCMAKE_CXX_COMPILER_LAUNCHER=ccache",
             "-DBA_APPLE_BUNDLE=ON",
             "-DCMAKE_OSX_DEPLOYMENT_TARGET=11",
             "-DBA_PY_PLATFORM=" + dirs.python,
             "-B", dirs.build, "-S", dirs.source])

        return conf_prc

    def build(dirs:BuildDirs):
        """ Build in the given build directory """
        build_prc = subp.run(["cmake", "--build", dirs.build, "--config Release",
                              "--parallel %i" % MacMake_x64.NPROC])
        return build_prc

    def test(dirs:BuildDirs):
        """ Perform tests """
        ENV["NOPLOT"] = "ON"
        os.chdir(dirs.build)
        test_prc = subp.run(["ctest", "-C", "Release", "--parallel %i" % MacMake_x64.NPROC,
                             "--output-on-failure"])
        return test_prc

    def pack(dirs:BuildDirs):
        """ Build package """
        wheel_prc = subp.run(["zsh", dirs.build + "/var/mk_pypack_macos.zsh"])
        wheel_prc.check_returncode()
        pack_prc = subp.run(["cpack", dirs.build, "-B " + dirs.package])
        return pack_prc


class MacMake_arm:
    PLATFORM = "darwin"
    ARCH = "arm64"
    QT_DIR = "/opt/homebrew/opt/qt"
    QTCMake_DIR = QT_DIR + "lib/cmake"
    OPTDIR="/Users/Shared/Software"
    # nr of logical CPUs can be obtained via `sysctl hw.logicalcpu`
    NPROC = 14
    PY_PLATFORM_BASE = "/opt/multipython/Python"
    # supported Python platforms: default (3.9), 3.8, 3.10, 3.11
    SUPPORTED_PY_VERSIONS = ("", "38", "310", "311")

    def configure(dirs:BuildDirs):
        """
        Configure the build with a given source and build directory
        and a given Python platform
        """
        conf_prc = subp.run(
            ["cmake",
             "-DCMAKE_PREFIX_PATH=" + MacMake_arm.OPTDIR + ";" + MacMake_arm.QTCMake_DIR,
             "-DCMAKE_BUILD_TYPE=Release",
             "-DCMAKE_CXX_COMPILER_LAUNCHER=ccache",
             "-DBA_APPLE_BUNDLE=ON",
             "-DBA_PY_PLATFORM=" + dirs.python,
             "-B", dirs.build, "-S", dirs.source])

        return conf_prc

    def build(dirs:BuildDirs):
        """ Build in the given build directory """
        build_prc = subp.run(["cmake", "--build", dirs.build, "--config Release",
                              "--parallel %i" % MacMake_arm.NPROC])
        return build_prc

    def test(dirs:BuildDirs):
        """ Perform tests """
        ENV["NOPLOT"] = "ON"
        os.chdir(dirs.build)
        test_prc = subp.run(["ctest", "-C", "Release", "--parallel %i" % MacMake_arm.NPROC,
                             "--output-on-failure"])
        return test_prc

    def pack(dirs:BuildDirs):
        """ Build package """
        wheel_prc = subp.run(["zsh", dirs.build + "/var/mk_pypack_macos.zsh"])
        wheel_prc.check_returncode()
        pack_prc = subp.run(["cpack", dirs.build, "-B " + dirs.package])
        return pack_prc


def make(maker, source_dir, python_version):
    """ Make for the given version of Python """
    print("#--- DIAGNOSTICS ---")
    # list all environmental variables
    for key, val in ENV.items():
        print(key, "=", val)

    print(SEPLINE)

    # set build-related directories
    source_dir = os.path.abspath(source_dir)
    build_dir = source_dir + "/" + BUILD_DIR
    package_dir = "installer"
    python_platform = ""
    python_version = python_version.strip()
    if python_version:
        build_dir = build_dir + "_py" + python_version
        package_dir = package_dir + "_py" + python_version
        python_platform = maker.PY_PLATFORM_BASE + python_version

    dirs = BuildDirs(source_dir, build_dir,
                     build_dir + "/" + package_dir, python_platform)

    print("* build directories:", dirs)

    # make the CMake build directory and switch to it
    if os.path.exists(dirs.build):
        raise RuntimeError("Build directory '%s' already exists" % dirs.build)

    os.mkdir(dirs.build)
    os.chdir(dirs.build)

    print("\n#--- CONFIGURE ---")
    print(subp.check_output(["cmake", "--version"], text=True, encoding="utf-8"))
    print("current path:" , os.getcwd())
    conf_prc = maker.configure(dirs)
    # NOTE: check_returncode() raises a `CalledProcessError` exception if the process fails
    conf_prc.check_returncode()

    print("\n#--- BUILD ---")
    build_start = timer()
    build_prc = maker.build(dirs)
    build_prc.check_returncode()
    build_elapsed = timer() - build_start
    print(SEPLINE)

    print("\n#--- TEST ---")
    test_start = timer()
    test_prc = maker.test(dirs)
    test_prc.check_returncode()
    test_elapsed = timer() - test_start
    print(SEPLINE)

    print("\n#--- PACKAGE ---")
    print("package directory: '%s'" % package_dir)
    pack_start = timer()
    pack_prc = maker.pack(dirs)
    pack_prc.check_returncode()
    pack_elapsed = timer() - pack_start
    print(SEPLINE)

    print()
    print("#--- Total build time = {:.3f} sec".format(build_elapsed))
    print("#--- Total test time = {:.3f} sec".format(test_elapsed))
    print("#--- Total packaging time = {:.3f} sec".format(pack_elapsed))
    print()


def make_all(source_dir="."):
    """ Make for all versions of Python """
    # supported platforms; see `platform.system` documentation
    SUPPORTED_PLATFORMS = (WindowsMake.PLATFORM, MacMake_x64.PLATFORM,
                           MacMake_arm.PLATFORM, LinuxMake.PLATFORM)

    # switch to the proper make commands for the platform
    platform_ = platform.system().lower()
    arch_ = platform.machine().lower()

    if not platform_ in SUPPORTED_PLATFORMS:
        raise NotImplementedError(
            "MuliPythonBuild: Platform '%s' is not in the supported platforms %s"
            % (platform_, SUPPORTED_PLATFORMS))

    if platform_ == LinuxMake.PLATFORM and arch_ == LinuxMake.ARCH:
        maker = LinuxMake
    elif platform_ == WindowsMake.PLATFORM and arch_ == WindowsMake.ARCH:
        maker = WindowsMake
    elif platform_ == MacMake_x64.PLATFORM:
        if arch_ == MacMake_x64.ARCH:
            maker = MacMake_x64
        elif arch_ == MacMake_arm.ARCH:
            maker = MacMake_arm
        else:
            raise NotImplementedError(
                "MuliPythonBuild: Architecture '%s' on platform '%s' is not supported"
                % (arch_, platform_))
    else:
        raise NotImplementedError(
            "MuliPythonBuild: Platform '%s' or architecture '%s' is not supported"
            % (platform_, arch_))

    source_dir = os.path.abspath(source_dir)
    for python_version in maker.SUPPORTED_PY_VERSIONS:
        py_ver_str = python_version if python_version else "default"
        print(SEPLINE)
        print("*** Build for Python version '%s' ***" % py_ver_str)
        make(maker, source_dir, python_version)
        print("*** END: Build for Python version '%s' ***" % py_ver_str)
        print(SEPLINE)


if __name__ == "__main__":
    make_all()
