#!/usr/bin/env python3
#
# Licensed to the .NET Foundation under one or more agreements.
# The .NET Foundation licenses this file to you under the MIT license.
#
# Title               : superpmi_collect_setup.py
#
# Notes:
#
# Script to setup directory structure required to perform SuperPMI collection in CI.
# It does the following steps:
# 1.  It creates `correlation_payload_directory` that contains files from CORE_ROOT, src\coreclr\scripts.
#     This directory is the one that is sent to all the helix machines that performs SPMI collection.
# 2.  It clones dotnet/jitutils, builds it and then copies the `pmi.dll` to `correlation_payload_directory` folder.
#     This file is needed to do pmi SPMI runs.
# 3.  The script takes `input_artifacts` parameter which contains managed .dlls and .exes on
#     which SPMI needs to be run. This script will partition these folders into equal buckets of approximately `max_size`
#     bytes and stores them under `payload` directory. Each sub-folder inside `payload` directory is sent to individual
#     helix machine to do SPMI collection on. E.g. for `input_artifacts` to be run on libraries, the parameter would be path to
#     `CORE_ROOT` folder and this script will copy `max_size` bytes of those files under `payload/libraries/0/binaries`,
#     `payload/libraries/1/binaries` and so forth.
# 4.  Lastly, it sets the pipeline variables.
#
# Below are the helix queues it sets depending on the OS/architecture:
# | Arch  | windows                 | Linux                                                                                                                                |
# |-------|-------------------------|--------------------------------------------------------------------------------------------------------------------------------------|
# | x86   | Windows.10.Amd64.X86.Rt |                                                                                                                                      |
# | x64   | Windows.10.Amd64.X86.Rt | Ubuntu.1804.Amd64                                                                                                                    |
# | arm   | -                       | (Ubuntu.1804.Arm32)Ubuntu.1804.Armarch@mcr.microsoft.com/dotnet-buildtools/prereqs:ubuntu-18.04-helix-arm32v7-bfcd90a-20200121150440 |
# | arm64 | Windows.10.Arm64        | (Ubuntu.1804.Arm64)Ubuntu.1804.ArmArch@mcr.microsoft.com/dotnet-buildtools/prereqs:ubuntu-18.04-helix-arm64v8-20210531091519-97d8652 |
#
################################################################################
################################################################################

import argparse
import os
import stat

from coreclr_arguments import *
from jitutil import run_command, copy_directory, copy_files, set_pipeline_variable, ChangeDir, TempDir


# Start of parser object creation.

parser = argparse.ArgumentParser(description="description")

parser.add_argument("-source_directory", help="path to source directory")
parser.add_argument("-core_root_directory", help="path to core_root directory")
parser.add_argument("-arch", help="Architecture")
parser.add_argument("-mch_file_tag", help="Tag to be used to mch files")
parser.add_argument("-collection_name", help="Name of the SPMI collection to be done (e.g., libraries, tests)")
parser.add_argument("-collection_type", help="Type of the SPMI collection to be done (crossgen, crossgen2, pmi)")
parser.add_argument("-input_directory", help="directory containing assemblies for which superpmi collection to be done")
parser.add_argument("-max_size", help="Max size of each partition in MB")
is_windows = platform.system() == "Windows"

native_binaries_to_ignore = [
    "api-ms-win-core-console-l1-1-0.dll",
    "api-ms-win-core-datetime-l1-1-0.dll",
    "api-ms-win-core-debug-l1-1-0.dll",
    "api-ms-win-core-errorhandling-l1-1-0.dll",
    "api-ms-win-core-file-l1-1-0.dll",
    "api-ms-win-core-file-l1-2-0.dll",
    "api-ms-win-core-file-l2-1-0.dll",
    "api-ms-win-core-handle-l1-1-0.dll",
    "api-ms-win-core-heap-l1-1-0.dll",
    "api-ms-win-core-interlocked-l1-1-0.dll",
    "api-ms-win-core-libraryloader-l1-1-0.dll",
    "api-ms-win-core-localization-l1-2-0.dll",
    "api-ms-win-core-memory-l1-1-0.dll",
    "api-ms-win-core-namedpipe-l1-1-0.dll",
    "api-ms-win-core-processenvironment-l1-1-0.dll",
    "api-ms-win-core-processthreads-l1-1-0.dll",
    "api-ms-win-core-processthreads-l1-1-1.dll",
    "api-ms-win-core-profile-l1-1-0.dll",
    "api-ms-win-core-rtlsupport-l1-1-0.dll",
    "api-ms-win-core-string-l1-1-0.dll",
    "api-ms-win-core-synch-l1-1-0.dll",
    "api-ms-win-core-synch-l1-2-0.dll",
    "api-ms-win-core-sysinfo-l1-1-0.dll",
    "api-ms-win-core-timezone-l1-1-0.dll",
    "api-ms-win-core-util-l1-1-0.dll",
    "api-ms-win-crt-conio-l1-1-0.dll",
    "api-ms-win-crt-convert-l1-1-0.dll",
    "api-ms-win-crt-environment-l1-1-0.dll",
    "api-ms-win-crt-filesystem-l1-1-0.dll",
    "api-ms-win-crt-heap-l1-1-0.dll",
    "api-ms-win-crt-locale-l1-1-0.dll",
    "api-ms-win-crt-math-l1-1-0.dll",
    "api-ms-win-crt-multibyte-l1-1-0.dll",
    "api-ms-win-crt-private-l1-1-0.dll",
    "api-ms-win-crt-process-l1-1-0.dll",
    "api-ms-win-crt-runtime-l1-1-0.dll",
    "api-ms-win-crt-stdio-l1-1-0.dll",
    "api-ms-win-crt-string-l1-1-0.dll",
    "api-ms-win-crt-time-l1-1-0.dll",
    "api-ms-win-crt-utility-l1-1-0.dll",
    "clretwrc.dll",
    "clrgc.dll",
    "clrjit.dll",
    "clrjit_unix_arm_arm.dll",
    "clrjit_unix_arm_arm64.dll",
    "clrjit_unix_arm_x64.dll",
    "clrjit_unix_arm_x86.dll",
    "clrjit_unix_arm64_arm64.dll",
    "clrjit_unix_arm64_x64.dll",
    "clrjit_unix_armel_arm.dll",
    "clrjit_unix_armel_arm64.dll",
    "clrjit_unix_armel_x64.dll",
    "clrjit_unix_armel_x86.dll",
    "clrjit_unix_osx_arm64_arm64.dll",
    "clrjit_unix_osx_arm64_x64.dll",
    "clrjit_unix_x64_arm64.dll",
    "clrjit_unix_x64_x64.dll",
    "clrjit_win_arm_arm.dll",
    "clrjit_win_arm_arm64.dll",
    "clrjit_win_arm_x64.dll",
    "clrjit_win_arm_x86.dll",
    "clrjit_win_arm64_arm64.dll",
    "clrjit_win_arm64_x64.dll",
    "clrjit_win_x64_arm64.dll",
    "clrjit_win_x64_x64.dll",
    "clrjit_win_x86_arm.dll",
    "clrjit_win_x86_arm64.dll",
    "clrjit_win_x86_x64.dll",
    "clrjit_win_x86_x86.dll",
    "clrjit_universal_arm_arm.dll",
    "clrjit_universal_arm_arm64.dll",
    "clrjit_universal_arm_x64.dll",
    "clrjit_universal_arm_x86.dll",
    "clrjit_universal_arm64_arm64.dll",
    "clrjit_universal_arm64_x64.dll",
    "coreclr.dll",
    "CoreConsole.exe",
    "coredistools.dll",
    "CoreRun.exe",
    "CoreShim.dll",
    "createdump.exe",
    "crossgen.exe",
    "crossgen2.exe",
    "dbgshim.dll",
    "ilasm.exe",
    "ildasm.exe",
    "jitinterface_arm.dll",
    "jitinterface_arm64.dll",
    "jitinterface_x64.dll",
    "jitinterface_x86.dll",
    "KernelTraceControl.dll",
    "KernelTraceControl.Win61.dll",
    "mcs.exe",
    "Microsoft.DiaSymReader.Native.amd64.dll",
    "Microsoft.DiaSymReader.Native.x86.dll",
    "mscordaccore.dll",
    "mscordbi.dll",
    "mscorrc.dll",
    "msdia140.dll",
    "R2RDump.exe",
    "R2RTest.exe",
    "superpmi.exe",
    "superpmi-shim-collector.dll",
    "superpmi-shim-counter.dll",
    "superpmi-shim-simple.dll",
    "System.IO.Compression.Native.dll",
    "ucrtbase.dll",
    "xunit.console.exe",
]

MAX_FILES_COUNT = 1500


def setup_args(args):
    """ Setup the args for SuperPMI to use.

    Args:
        args (ArgParse): args parsed by arg parser

    Returns:
        args (CoreclrArguments)

    """
    coreclr_args = CoreclrArguments(args, require_built_core_root=False, require_built_product_dir=False,
                                    require_built_test_dir=False, default_build_type="Checked")

    coreclr_args.verify(args,
                        "source_directory",
                        lambda source_directory: os.path.isdir(source_directory),
                        "source_directory doesn't exist")

    coreclr_args.verify(args,
                        "core_root_directory",
                        lambda core_root_directory: os.path.isdir(core_root_directory),
                        "core_root_directory doesn't exist")

    coreclr_args.verify(args,
                        "arch",
                        lambda unused: True,
                        "Unable to set arch")

    coreclr_args.verify(args,
                        "mch_file_tag",
                        lambda unused: True,
                        "Unable to set mch_file_tag")

    coreclr_args.verify(args,
                        "collection_name",
                        lambda unused: True,
                        "Unable to set collection_name")

    coreclr_args.verify(args,
                        "collection_type",
                        lambda unused: True,
                        "Unable to set collection_type")

    coreclr_args.verify(args,
                        "input_directory",
                        lambda input_directory: os.path.isdir(input_directory),
                        "input_directory doesn't exist")

    coreclr_args.verify(args,
                        "max_size",
                        lambda max_size: max_size > 0,
                        "Please enter valid positive numeric max_size",
                        modify_arg=lambda max_size: int(
                            max_size) * 1000 * 1000 if max_size is not None and max_size.isnumeric() else 0
                        # Convert to MB
                        )
    return coreclr_args


def get_files_sorted_by_size(src_directory, exclude_directories, exclude_files):
    """ For a given src_directory, returns all the .dll files sorted by size.

    Args:
        src_directory (string): Path of directory to enumerate.
        exclude_directories ([string]): Directory names to exclude.
        exclude_files ([string]): File names to exclude.
    """

    def sorter_by_size(pair):
        """ Sorts the pair (file_name, file_size) tuple in descending order of file_size

        Args:
            pair ([(string, int)]): List of tuple of file_name, file_size
        """
        pair.sort(key=lambda x: x[1], reverse=True)
        return pair

    filename_with_size = []

    for file_path, dirs, files in os.walk(src_directory, topdown=True):
        # Credit: https://stackoverflow.com/a/19859907
        dirs[:] = [d for d in dirs if d not in exclude_directories]
        for name in files:
            # Make the exclude check case-insensitive
            exclude_files_lower = [filename.lower() for filename in exclude_files]
            if name.lower() in exclude_files_lower:
                continue
            curr_file_path = os.path.join(file_path, name)

            if not os.path.isfile(curr_file_path):
                continue
            if not name.endswith(".dll") and not name.endswith(".exe"):
                continue

            size = os.path.getsize(curr_file_path)
            filename_with_size.append((curr_file_path, size))

    return sorter_by_size(filename_with_size)


def first_fit(sorted_by_size, max_size):
    """ Given a list of file names along with size in descending order, divides the files
    in number of buckets such that each bucket doesn't exceed max_size. Since this is a first-fit
    approach, it doesn't guarantee to find the bucket with tighest spot available.

    Args:
        sorted_by_size ((string, int)): (file_name, file_size) tuple
        max_size (int): Maximum size (in bytes) of each bucket.

    Returns:
        [{int, [string]}]: Returns a dictionary of partition-index to list of file names following in that bucket.
    """
    partitions = {}
    for curr_file in sorted_by_size:
        _, file_size = curr_file

        # Find the right bucket
        found_bucket = False

        if file_size < max_size:
            for p_index in partitions:
                total_in_curr_par = sum(n for _, n in partitions[p_index])
                if ((total_in_curr_par + file_size) < max_size) and (len(partitions[p_index]) < MAX_FILES_COUNT):
                    partitions[p_index].append(curr_file)
                    found_bucket = True
                    break

            if not found_bucket:
                partitions[len(partitions)] = [curr_file]

    total_size = 0
    for p_index in partitions:
        partition_size = sum(n for _, n in partitions[p_index])
        print("Partition {0}: {1} files with {2} bytes.".format(p_index, len(partitions[p_index]), partition_size))
        total_size += partition_size
    print("Total {0} partitions with {1} bytes.".format(str(len(partitions)), total_size))

    return partitions


def partition_files(src_directory, dst_directory, max_size, exclude_directories=[],
                    exclude_files=native_binaries_to_ignore):
    """ Copy bucketized files based on size to destination folder.

    Args:
        src_directory (string): Source folder containing files to be copied.
        dst_directory (string): Destination folder where files should be copied.
        max_size (int): Maximum partition size in bytes
        exclude_directories ([string]): List of folder names to be excluded.
        exclude_files ([string]): List of files names to be excluded.
    """

    print('Partitioning files from {0} to {1}'.format(src_directory, dst_directory))
    sorted_by_size = get_files_sorted_by_size(src_directory, exclude_directories, exclude_files)
    partitions = first_fit(sorted_by_size, max_size)

    index = 0
    for p_index in partitions:
        file_names = [curr_file[0] for curr_file in partitions[p_index]]
        curr_dst_path = os.path.join(dst_directory, str(index), "binaries")
        copy_files(src_directory, curr_dst_path, file_names)
        index += 1


def setup_microbenchmark(workitem_directory, arch):
    """ Perform setup of microbenchmarks

    Args:
        workitem_directory (string): Path to work
        arch (string): Architecture for which dotnet will be installed
    """
    performance_directory = os.path.join(workitem_directory, "performance")

    run_command(
        ["git", "clone", "--quiet", "--depth", "1", "https://github.com/dotnet/performance", performance_directory])

    with ChangeDir(performance_directory):
        dotnet_directory = os.path.join(performance_directory, "tools", "dotnet", arch)
        dotnet_install_script = os.path.join(performance_directory, "scripts", "dotnet.py")

        if not os.path.isfile(dotnet_install_script):
            print("Missing " + dotnet_install_script)
            return

        run_command(
            get_python_name() + [dotnet_install_script, "install", "--architecture", arch, "--install-dir",
                                 dotnet_directory, "--verbose"])


def get_python_name():
    """Gets the python name

    Returns:
        [string]: Returns the appropriate python name depending on the OS.
    """
    if is_windows:
        return ["py", "-3"]
    else:
        return ["python3"]


def main(main_args):
    """ Main entrypoint

    Args:
        main_args ([type]): Arguments to the script
    """
    coreclr_args = setup_args(main_args)
    source_directory = coreclr_args.source_directory

    # CorrelationPayload directories
    correlation_payload_directory = os.path.join(coreclr_args.source_directory, "payload")
    superpmi_src_directory = os.path.join(source_directory, 'src', 'coreclr', 'scripts')
    superpmi_dst_directory = os.path.join(correlation_payload_directory, "superpmi")
    arch = coreclr_args.arch
    helix_source_prefix = "official"
    creator = ""
    ci = True
    if is_windows:
        helix_queue = "Windows.10.Arm64" if arch == "arm64" else "Windows.10.Amd64.X86.Rt"
    else:
        if arch == "arm":
            helix_queue = "(Ubuntu.1804.Arm32)Ubuntu.1804.Armarch@mcr.microsoft.com/dotnet-buildtools/prereqs:ubuntu-18.04-helix-arm32v7-bfcd90a-20200121150440"
        elif arch == "arm64":
            helix_queue = "(Ubuntu.1804.Arm64)Ubuntu.1804.ArmArch@mcr.microsoft.com/dotnet-buildtools/prereqs:ubuntu-18.04-helix-arm64v8-20210531091519-97d8652"
        else:
            helix_queue = "Ubuntu.1804.Amd64"

    # create superpmi directory
    print('Copying {} -> {}'.format(superpmi_src_directory, superpmi_dst_directory))
    copy_directory(superpmi_src_directory, superpmi_dst_directory, verbose_output=True, match_func=lambda path: any(path.endswith(extension) for extension in [".py"]))

    if is_windows:
        acceptable_copy = lambda path: any(path.endswith(extension) for extension in [".py", ".dll", ".exe", ".json"])
    else:
        # Need to accept files without any extension, which is how executable file's names look.
        acceptable_copy = lambda path: (os.path.basename(path).find(".") == -1) or any(path.endswith(extension) for extension in [".py", ".dll", ".so", ".json"])

    print('Copying {} -> {}'.format(coreclr_args.core_root_directory, superpmi_dst_directory))
    copy_directory(coreclr_args.core_root_directory, superpmi_dst_directory, verbose_output=True, match_func=acceptable_copy)

    # Copy all the test files to CORE_ROOT
    # The reason is there are lot of dependencies with *.Tests.dll and to ensure we do not get
    # Reflection errors, just copy everything to CORE_ROOT so for all individual partitions, the
    # references will be present in CORE_ROOT.
    if coreclr_args.collection_name == "libraries_tests":
        print('Copying {} -> {}'.format(coreclr_args.input_directory, superpmi_dst_directory))

        def make_readable(folder_name):
            """Make file executable by changing the permission

            Args:
                folder_name (string): folder to mark with 744
            """
            if is_windows:
                return

            print("Inside make_readable")
            run_command(["ls", "-l", folder_name])
            for file_path, dirs, files in os.walk(folder_name, topdown=True):
                for d in dirs:
                    os.chmod(os.path.join(file_path, d),
                    # read+write+execute for owner
                    (stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR) |
                    # read for group
                    (stat.S_IRGRP) |
                    # read for other
                    (stat.S_IROTH))

                for f in files:
                    os.chmod(os.path.join(file_path, f),
                    # read+write+execute for owner
                    (stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR) |
                    # read for group
                    (stat.S_IRGRP) |
                    # read for other
                    (stat.S_IROTH))
            run_command(["ls", "-l", folder_name])

        make_readable(coreclr_args.input_directory)
        copy_directory(coreclr_args.input_directory, superpmi_dst_directory, verbose_output=True, match_func=acceptable_copy)

    # Workitem directories
    workitem_directory = os.path.join(source_directory, "workitem")
    input_artifacts = ""

    if coreclr_args.collection_name == "benchmarks":
        # Setup microbenchmarks
        setup_microbenchmark(workitem_directory, arch)
    else:
        # Setup for pmi/crossgen runs

        # Clone and build jitutils
        try:
            with TempDir() as jitutils_directory:
                run_command(
                    ["git", "clone", "--quiet", "--depth", "1", "https://github.com/dotnet/jitutils", jitutils_directory])

                # Make sure ".dotnet" directory exists, by running the script at least once
                dotnet_script_name = "dotnet.cmd" if is_windows else "dotnet.sh"
                dotnet_script_path = os.path.join(source_directory, dotnet_script_name)
                run_command([dotnet_script_path, "--info"], jitutils_directory)

                # Set dotnet path to run build
                os.environ["PATH"] = os.path.join(source_directory, ".dotnet") + os.pathsep + os.environ["PATH"]
                build_file = "build.cmd" if is_windows else "build.sh"
                run_command([os.path.join(jitutils_directory, build_file), "-p"], jitutils_directory)

                copy_files(os.path.join(jitutils_directory, "bin"), superpmi_dst_directory, [os.path.join(jitutils_directory, "bin", "pmi.dll")])
        except PermissionError as pe_error:
            # Details: https://bugs.python.org/issue26660
            print('Ignoring PermissionError: {0}'.format(pe_error))

        # NOTE: we can't use the build machine ".dotnet" to run on all platforms. E.g., the Windows x86 build uses a
        # Windows x64 .dotnet\dotnet.exe that can't load a 32-bit shim. Thus, we always use corerun from Core_Root to invoke crossgen2.
        # The following will copy .dotnet to the correlation payload in case we change our mind, and need or want to use it for some scenarios.

        # # Copy ".dotnet" to correlation_payload_directory for crossgen2 job; it is needed to invoke crossgen2.dll
        # if coreclr_args.collection_type == "crossgen2":
        #     dotnet_src_directory = os.path.join(source_directory, ".dotnet")
        #     dotnet_dst_directory = os.path.join(correlation_payload_directory, ".dotnet")
        #     print('Copying {} -> {}'.format(dotnet_src_directory, dotnet_dst_directory))
        #     copy_directory(dotnet_src_directory, dotnet_dst_directory, verbose_output=False)

        # payload
        pmiassemblies_directory = os.path.join(workitem_directory, "pmiAssembliesDirectory")
        input_artifacts = os.path.join(pmiassemblies_directory, coreclr_args.collection_name)
        exclude_directory = ['Core_Root'] if coreclr_args.collection_name == "coreclr_tests" else []
        exclude_files = native_binaries_to_ignore
        if coreclr_args.collection_type == "crossgen2":
            print('Adding exclusions for crossgen2')
            # Currently, trying to crossgen2 R2RTest\Microsoft.Build.dll causes a pop-up failure, so exclude it.
            exclude_files += ["Microsoft.Build.dll"]

        if coreclr_args.collection_name == "libraries_tests":
            # libraries_tests artifacts contains files from core_root folder. Exclude them.
            core_root_dir = coreclr_args.core_root_directory
            exclude_files += [item for item in os.listdir(core_root_dir)
                              if os.path.isfile(os.path.join(core_root_dir, item)) and (item.endswith(".dll") or item.endswith(".exe"))]

        partition_files(coreclr_args.input_directory, input_artifacts, coreclr_args.max_size, exclude_directory,
                        exclude_files)

    # Set variables
    print('Setting pipeline variables:')
    set_pipeline_variable("CorrelationPayloadDirectory", correlation_payload_directory)
    set_pipeline_variable("WorkItemDirectory", workitem_directory)
    set_pipeline_variable("InputArtifacts", input_artifacts)
    set_pipeline_variable("Python", ' '.join(get_python_name()))
    set_pipeline_variable("Architecture", arch)
    set_pipeline_variable("Creator", creator)
    set_pipeline_variable("Queue", helix_queue)
    set_pipeline_variable("HelixSourcePrefix", helix_source_prefix)
    set_pipeline_variable("MchFileTag", coreclr_args.mch_file_tag)


################################################################################
# __main__
################################################################################

if __name__ == "__main__":
    args = parser.parse_args()
    sys.exit(main(args))
