"""
Build BornAgain with multiple versions of Python
"""

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_PACKAGE=ON",
             "-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 via NSIS """
        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",
             "-DBA_PY_PACKAGE=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_pypack_linux.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_PACKAGE=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_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_PACKAGE=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()
