# @file Edk2ToolsBuild.py
# Invocable class that builds the basetool c files.
#
# Supports VS2019, VS2022, and GCC
##
# Copyright (c) Microsoft Corporation
#
# SPDX-License-Identifier: BSD-2-Clause-Patent
##
import os
import sys
import logging
import argparse
import multiprocessing
import shutil
from edk2toolext import edk2_logging
from edk2toolext.environment import self_describing_environment, shell_environment
from edk2toolext.base_abstract_invocable import BaseAbstractInvocable
from edk2toollib.utility_functions import RunCmd, GetHostInfo
from edk2toollib.windows.locate_tools import QueryVcVariables


class Edk2ToolsBuild(BaseAbstractInvocable):

    def ParseCommandLineOptions(self):
        ''' parse arguments '''
        ParserObj = argparse.ArgumentParser()
        ParserObj.add_argument("-t", "--tool_chain_tag", dest="tct", default="VS2022",
                               help="Set the toolchain used to compile the build tools")
        ParserObj.add_argument("-a", "--target_arch", dest="arch", default=None, choices=[None, 'IA32', 'X64', 'AARCH64'],
                               help="Specify the architecture of the built base tools. Not specifying this will fall back to the default "
                               "behavior, for Windows builds, IA32 target will be built, for Linux builds, target arch will be the same as host arch.")
        args = ParserObj.parse_args()
        self.tool_chain_tag = args.tct
        self.target_arch = args.arch

    def GetWorkspaceRoot(self):
        ''' Return the workspace root for initializing the SDE '''

        # this is the bastools dir...not the traditional EDK2 workspace root
        return os.path.dirname(os.path.abspath(__file__))

    def GetActiveScopes(self):
        ''' return tuple containing scopes that should be active for this process '''

        # Adding scope for cross compilers when building for AARCH64
        scopes = ('global',)
        if GetHostInfo().os == "Linux" and self.tool_chain_tag.lower().startswith("gcc"):
            if self.target_arch is None:
                return scopes
            if "AARCH64" in self.target_arch:
                scopes += ("gcc_aarch64_linux",)
        return scopes

    def GetLoggingLevel(self, loggerType):
        ''' Get the logging level for a given type (return Logging.Level)
        base == lowest logging level supported
        con  == Screen logging
        txt  == plain text file logging
        md   == markdown file logging
        '''
        if(loggerType == "con"):
            return logging.ERROR
        else:
            return logging.DEBUG

    def GetLoggingFolderRelativeToRoot(self):
        ''' Return a path to folder for log files '''
        return "BaseToolsBuild"

    def GetVerifyCheckRequired(self):
        ''' Will call self_describing_environment.VerifyEnvironment if this returns True '''
        return True

    def GetLoggingFileName(self, loggerType):
        ''' Get the logging file name for the type.
        Return None if the logger shouldn't be created

        base == lowest logging level supported
        con  == Screen logging
        txt  == plain text file logging
        md   == markdown file logging
        '''
        return "BASETOOLS_BUILD"

    def WritePathEnvFile(self, OutputDir):
        ''' Write a PyTool path env file for future PyTool based edk2 builds'''
        content = '''##
# Set shell variable EDK_TOOLS_BIN to this folder
#
# Autogenerated by Edk2ToolsBuild.py
#
# Copyright (c), Microsoft Corporation
# SPDX-License-Identifier: BSD-2-Clause-Patent
##
{
  "id": "You-Built-BaseTools",
  "scope": "edk2-build",
  "flags": ["set_shell_var", "set_path"],
  "var_name": "EDK_TOOLS_BIN"
}
'''
        with open(os.path.join(OutputDir, "basetoolsbin_path_env.yaml"), "w") as f:
            f.write(content)

    def CleanBuildOutputs(self, shell_env, make_cmd, make_clean_target):
        '''Clean all build outputs.'''
        base_tools_path = shell_env.get_shell_var("EDK_TOOLS_PATH")

        # Remove the BaseToolsBuild log/temp folder.
        # Close all file-based log handlers first so the log file is not locked.
        build_folder = os.path.join(base_tools_path, self.GetLoggingFolderRelativeToRoot())
        if os.path.isdir(build_folder):
            logging.info("Removing directory: %s", build_folder)
            root_logger = logging.getLogger()
            for handler in list(root_logger.handlers):
                if isinstance(handler, logging.FileHandler):
                    handler.close()
                    root_logger.removeHandler(handler)
            shutil.rmtree(build_folder)

        # Re-create the log folder and file handler so subsequent build output is logged.
        os.makedirs(build_folder, exist_ok=True)
        log_file = os.path.join(build_folder, self.GetLoggingFileName("txt") + ".txt")
        file_handler = logging.FileHandler(log_file, mode='w')
        file_handler.setLevel(logging.DEBUG)
        logging.getLogger().addHandler(file_handler)
        logging.info("Running Python version: " + str(sys.version_info))

        # Do the actual clean
        ret = RunCmd(make_cmd, make_clean_target, workingdir=base_tools_path)
        if ret != 0:
            logging.warning("%s %s returned %d", make_cmd, make_clean_target, ret)

        # Remove output bin directories to delete stale YAML files.
        for bin_dir in [os.path.join(base_tools_path, "Bin", "Win32"),
                        os.path.join(base_tools_path, "Bin", "Win64"),
                        os.path.join(base_tools_path, "Source", "C", "bin")]:
            if os.path.isdir(bin_dir):
                logging.info("Removing directory: %s", bin_dir)
                shutil.rmtree(bin_dir)

        logging.info("Clean complete.")

    def Go(self):
        logging.info("Running Python version: " + str(sys.version_info))

        (build_env, shell_env) = self_describing_environment.BootstrapEnvironment(
            self.GetWorkspaceRoot(), self.GetActiveScopes())

        # # Bind our current execution environment into the shell vars.
        ph = os.path.dirname(sys.executable)
        if " " in ph:
            ph = '"' + ph + '"'
        shell_env.set_shell_var("PYTHON_HOME", ph)
        # PYTHON_COMMAND is required to be set for using edk2 python builds.
        pc = sys.executable
        if " " in pc:
            pc = '"' + pc + '"'
        shell_env.set_shell_var("PYTHON_COMMAND", pc)

        if self.tool_chain_tag.lower().startswith("vs"):
            if self.target_arch is None:
                # Put a default as IA32
                self.target_arch = "IA32"

            if self.target_arch == "IA32":
                VcToolChainArch = "x86"
                TargetInfoArch = "x86"
                OutputDir = "Win32"
            elif self.target_arch == "X64":
                VcToolChainArch = "amd64"
                TargetInfoArch = "x86"
                OutputDir = "Win64"
            elif self.target_arch == "AARCH64":
                VcToolChainArch = "amd64_arm64"
                TargetInfoArch = "ARM"
                OutputDir = "Win64"
            else:
                raise NotImplementedError()

            self.OutputDir = os.path.join(
                shell_env.get_shell_var("EDK_TOOLS_PATH"), "Bin", OutputDir)

            # compiled tools need to be added to path because antlr is referenced
            HostInfo = GetHostInfo()
            if TargetInfoArch == HostInfo.arch:
                # not cross compiling
                shell_env.insert_path(self.OutputDir)
            else:
                # cross compiling:
                # as the VfrCompile tool is needed in the build process, we need
                # to build one for the host system, then add the path to the
                # tools to the PATH environment variable
                shell_environment.CheckpointBuildVars()
                if HostInfo.arch == "x86" and HostInfo.bit == "64":
                    host_arch = "X64"
                    host_toolchain_arch = "amd64"
                    TempOutputDir = os.path.join(shell_env.get_shell_var("EDK_TOOLS_PATH"), "Bin", "Win64")
                elif HostInfo.arch == "x86" and HostInfo.bit == "32":
                    host_arch = "IA32"
                    host_toolchain_arch = "x86"
                    TempOutputDir = os.path.join(shell_env.get_shell_var("EDK_TOOLS_PATH"), "Bin", "Win32")
                elif HostInfo.arch == "ARM" and HostInfo.bit == "64":
                    host_arch = "AARCH64"
                    host_toolchain_arch = "amd64_arm64"
                    TempOutputDir = os.path.join(shell_env.get_shell_var("EDK_TOOLS_PATH"), "Bin", "Win64")
                else:
                    raise Exception("Unsupported host system. %s %s" % (HostInfo.arch, HostInfo.bit))

                interesting_keys = ["ExtensionSdkDir", "INCLUDE", "LIB"]
                interesting_keys.extend(
                    ["LIBPATH", "Path", "UniversalCRTSdkDir", "UCRTVersion", "WindowsLibPath", "WindowsSdkBinPath"])
                interesting_keys.extend(
                    ["WindowsSdkDir", "WindowsSdkVerBinPath", "WindowsSDKVersion", "VCToolsInstallDir"])
                vc_vars = QueryVcVariables(
                    interesting_keys, host_toolchain_arch, vs_version=self.tool_chain_tag.lower())
                for key in vc_vars.keys():
                    logging.debug(f"Var - {key} = {vc_vars[key]}")
                    if key.lower() == 'path':
                        shell_env.set_path(vc_vars[key])
                    else:
                        shell_env.set_shell_var(key, vc_vars[key])

                # Note: This HOST_ARCH is in respect to the BUILT base tools, not the host arch where
                # this script is BUILDING the base tools.
                shell_env.set_shell_var('HOST_ARCH', host_arch)
                shell_env.insert_path(TempOutputDir)

                # All set, build the tools for the host system.
                ret = RunCmd('nmake.exe', None,
                             workingdir=shell_env.get_shell_var("EDK_TOOLS_PATH"))
                if ret != 0:
                    raise Exception("Failed to build base tools for host system.")

                # Copy the output to a temp directory
                TempFolder = os.path.join(shell_env.get_shell_var("EDK_TOOLS_PATH"), "BaseToolsBuild", "Temp")
                if not os.path.exists(TempFolder):
                    os.makedirs(TempFolder)
                for file in os.listdir(TempOutputDir):
                    shutil.copy(os.path.join(TempOutputDir, file), TempFolder)

                # Clean up the build output
                ret = RunCmd('nmake.exe', 'cleanall',
                             workingdir=shell_env.get_shell_var("EDK_TOOLS_PATH"))

                # Remove the entire TempOutputDir
                shutil.rmtree(TempOutputDir)

                shell_environment.RevertBuildVars()
                shell_env.insert_path(TempFolder)

            # # Update environment with required VC vars.
            interesting_keys = ["ExtensionSdkDir", "INCLUDE", "LIB"]
            interesting_keys.extend(
                ["LIBPATH", "Path", "UniversalCRTSdkDir", "UCRTVersion", "WindowsLibPath", "WindowsSdkBinPath"])
            interesting_keys.extend(
                ["WindowsSdkDir", "WindowsSdkVerBinPath", "WindowsSDKVersion", "VCToolsInstallDir"])
            vc_vars = QueryVcVariables(
                interesting_keys, VcToolChainArch, vs_version=self.tool_chain_tag.lower())
            for key in vc_vars.keys():
                logging.debug(f"Var - {key} = {vc_vars[key]}")
                if key.lower() == 'path':
                    shell_env.set_path(vc_vars[key])
                else:
                    shell_env.set_shell_var(key, vc_vars[key])

            # Note: This HOST_ARCH is in respect to the BUILT base tools, not the host arch where
            # this script is BUILDING the base tools.
            shell_env.set_shell_var('HOST_ARCH', self.target_arch)

            self.CleanBuildOutputs(shell_env, 'nmake.exe', 'cleanall')

            # Actually build the tools.
            output_stream = edk2_logging.create_output_stream()
            ret = RunCmd('nmake.exe', None,
                         workingdir=shell_env.get_shell_var("EDK_TOOLS_PATH"))
            edk2_logging.remove_output_stream(output_stream)
            problems = edk2_logging.scan_compiler_output(output_stream)
            for level, problem in problems:
                logging.log(level, problem)
            if ret != 0:
                raise Exception("Failed to build.")

            self.WritePathEnvFile(self.OutputDir)
            return ret

        elif self.tool_chain_tag.lower().startswith("gcc") or self.tool_chain_tag.lower().startswith("clang"):
            # Note: This HOST_ARCH is in respect to the BUILT base tools, not the host arch where
            # this script is BUILDING the base tools.
            HostInfo = GetHostInfo()
            prefix = None
            TargetInfoArch = None
            if self.target_arch is not None:
                shell_env.set_shell_var('HOST_ARCH', self.target_arch)

                if "AARCH64" in self.target_arch:
                    prefix = shell_env.get_shell_var("GCC_AARCH64_PREFIX")
                    if prefix == None:
                        # now check for install dir.  If set then set the Prefix
                        install_path = shell_environment.GetEnvironment().get_shell_var("GCC_AARCH64_INSTALL")

                        # make GCC_AARCH64_PREFIX to align with tools_def.txt
                        prefix = os.path.join(install_path, "bin", "aarch64-none-linux-gnu-")

                    shell_environment.GetEnvironment().set_shell_var("GCC_PREFIX", prefix)
                    TargetInfoArch = "ARM"

                else:
                    TargetInfoArch = "x86"
            else:
                self.target_arch = HostInfo.arch
                TargetInfoArch = HostInfo.arch
            # Otherwise, the built binary arch will be consistent with the host system

            make_command = 'make'
            clang_bin = shell_env.get_shell_var("CLANG_BIN")
            if clang_bin is not None:
                if os.path.exists(os.path.join(clang_bin, "mingw32-make.exe")):
                    make_command = 'mingw32-make'

            # Added logic to support cross compilation scenarios
            if TargetInfoArch != HostInfo.arch:
                # this is defaulting to the version that comes with Ubuntu 20.04
                ver = shell_environment.GetBuildVars().GetValue("LIBUUID_VERSION", "2.34")
                work_dir = os.path.join(shell_env.get_shell_var("EDK_TOOLS_PATH"), self.GetLoggingFolderRelativeToRoot())
                pack_name = f"util-linux-{ver}"
                unzip_dir = os.path.join(work_dir, pack_name)

                if os.path.isfile(os.path.join(work_dir, f"{pack_name}.tar.gz")):
                    os.remove(os.path.join(work_dir, f"{pack_name}.tar.gz"))
                if os.path.isdir(unzip_dir):
                    shutil.rmtree(unzip_dir)

                # cross compiling, need to rebuild libuuid for the target
                ret = RunCmd("wget", f"https://mirrors.edge.kernel.org/pub/linux/utils/util-linux/v{ver}/{pack_name}.tar.gz", workingdir=work_dir)
                if ret != 0:
                    raise Exception(f"Failed to download libuuid version {ver} - {ret}")

                ret = RunCmd("tar", f"xvzf {pack_name}.tar.gz", workingdir=work_dir)
                if ret != 0:
                    raise Exception(f"Failed to untar the downloaded file {ret}")

                # configure the source to use the cross compiler
                pack_name = f"util-linux-{ver}"
                if "AARCH64" in self.target_arch:
                    ret = RunCmd("sh", f"./configure --host=aarch64-linux  -disable-all-programs --enable-libuuid CC={prefix}gcc", workingdir=unzip_dir)
                if ret != 0:
                    raise Exception(f"Failed to configure the util-linux to build with our gcc {ret}")

                ret = RunCmd(make_command, "", workingdir=unzip_dir)
                if ret != 0:
                    raise Exception(f"Failed to build the libuuid with our gcc {ret}")

                shell_environment.GetEnvironment().set_shell_var("CROSS_LIB_UUID", unzip_dir)
                shell_environment.GetEnvironment().set_shell_var("CROSS_LIB_UUID_INC", os.path.join(unzip_dir, "libuuid", "src"))

            self.CleanBuildOutputs(shell_env, make_command, 'clean')

            cpu_count = self.GetCpuThreads()

            output_stream = edk2_logging.create_output_stream()
            ret = RunCmd(make_command, f"-C .  -j {cpu_count}", workingdir=shell_env.get_shell_var("EDK_TOOLS_PATH"))
            edk2_logging.remove_output_stream(output_stream)
            problems = edk2_logging.scan_compiler_output(output_stream)
            for level, problem in problems:
                logging.log(level, problem)
            if ret != 0:
                raise Exception("Failed to build.")

            self.OutputDir = os.path.join(
                shell_env.get_shell_var("EDK_TOOLS_PATH"), "Source", "C", "bin")

            self.WritePathEnvFile(self.OutputDir)
            return ret

        logging.critical("Tool Chain not supported")
        return -1

    def GetCpuThreads(self) -> int:
        ''' Function to return number of cpus. If error return 1'''
        cpus = 1
        try:
            cpus = multiprocessing.cpu_count()
        except:
            # from the internet there are cases where cpu_count is not implemented.
            # will handle error by just doing single proc build
            pass
        return cpus



def main():
    Edk2ToolsBuild().Invoke()


if __name__ == "__main__":
    main()