# @file Edk2ToolsBuild.py
# Invocable class that builds the basetool c files.
#
# Supports VS2017, VS2019, and GCC5
##
# Copyright (c) Microsoft Corporation
#
# SPDX-License-Identifier: BSD-2-Clause-Patent
##
import os
import sys
import logging
import argparse
import multiprocessing
from edk2toolext import edk2_logging
from edk2toolext.environment import self_describing_environment
from edk2toolext.base_abstract_invocable import BaseAbstractInvocable
from edk2toollib.utility_functions import RunCmd
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="VS2017",
                               help="Set the toolchain used to compile the build tools")
        args = ParserObj.parse_args()
        self.tool_chain_tag = args.tct

    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 '''

        # for now don't use scopes
        return ('global',)

    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 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"):

            # # 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, 'x86', 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])

            self.OutputDir = os.path.join(
                shell_env.get_shell_var("EDK_TOOLS_PATH"), "Bin", "Win32")

            # compiled tools need to be added to path because antlr is referenced
            shell_env.insert_path(self.OutputDir)

            # 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"):
            cpu_count = self.GetCpuThreads()

            output_stream = edk2_logging.create_output_stream()
            ret = RunCmd("make", 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()