# Copyright (c) 2023, NVIDIA CORPORATION.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


"""Wrapper class to run tools associated with RAPIDS Accelerator for Apache Spark plugin on AWS-EMR."""
from spark_rapids_tools import CspEnv
from spark_rapids_pytools.cloud_api.sp_types import DeployMode
from spark_rapids_pytools.common.utilities import ToolLogging
from spark_rapids_pytools.rapids.bootstrap import Bootstrap
from spark_rapids_pytools.rapids.diagnostic import Diagnostic
from spark_rapids_pytools.rapids.qualification import QualFilterApp, QualificationAsLocal, \
    QualGpuClusterReshapeType
from spark_rapids_pytools.rapids.profiling import ProfilingAsLocal


class CliEmrLocalMode:  # pylint: disable=too-few-public-methods
    """
    A wrapper that runs the RAPIDS Accelerator tools locally on the dev machine.
    """

    @staticmethod
    def qualification(cpu_cluster: str = None,
                      eventlogs: str = None,
                      profile: str = None,
                      local_folder: str = None,
                      remote_folder: str = None,
                      gpu_cluster: str = None,
                      tools_jar: str = None,
                      filter_apps: str = QualFilterApp.tostring(QualFilterApp.SAVINGS),
                      gpu_cluster_recommendation: str = QualGpuClusterReshapeType.tostring(
                          QualGpuClusterReshapeType.get_default()),
                      jvm_heap_size: int = 24,
                      verbose: bool = False,
                      **rapids_options) -> None:
        """
        The Qualification tool analyzes Spark events generated from CPU based Spark applications to
        help quantify the expected acceleration and costs savings of migrating a Spark application
        or query to GPU. The wrapper downloads dependencies and executes the analysis on the local
        dev machine
        :param cpu_cluster: The EMR-cluster on which the Spark applications were executed. The argument
                can be an EMR-cluster or a valid path to the cluster's properties file (json format)
                generated by the AWS CLI
        :param  eventlogs: Event log filenames or S3 storage directories
                containing event logs (comma separated). If missing, the wrapper Reads the Spark's
                property `spark.eventLog.dir` defined in `cpu_cluster`. This property should be included
                in the output of `emr describe-cluster`. Note that the wrapper will raise an exception
                if the property is not set
        :param profile: A named AWS profile to get the settings/credentials of the AWS account
        :param local_folder: Local work-directory path to store the output and to be used as root
                directory for temporary folders/files. The final output will go into a subdirectory called
                ${local_folder}/qual-${EXEC_ID} where exec_id is an auto-generated unique identifier of the
                execution. If the argument is NONE, the default value is the env variable
                RAPIDS_USER_TOOLS_OUTPUT_DIRECTORY if any; or the current working directory.
        :param remote_folder: An S3 folder where the output is uploaded at the end of execution.
                If no value is provided, the output will be only available on local disk
        :param gpu_cluster: The EMR-cluster on which the Spark applications is planned to be migrated.
                The argument can be an EMR-cluster or a valid path to the cluster's properties file
                (json format) generated by the AWS CLI. If missing, the wrapper maps the EC2 machine
                instances of the original cluster into EC2 instances that support GPU acceleration
        :param tools_jar: Path to a bundled jar including Rapids tool. The path is a local filesystem,
                or remote S3 url. If missing, the wrapper downloads the latest rapids-4-spark-tools_*.jar
                from maven repo
        :param filter_apps: filtering criteria of the applications listed in the final STDOUT table
                is one of the following (ALL, SPEEDUPS, SAVINGS). Default is "SAVINGS".
                Note that this filter does not affect the CSV report.
                "ALL" means no filter applied. "SPEEDUPS" lists all the apps that are either
                'Recommended', or 'Strongly Recommended' based on speedups. "SAVINGS"
                lists all the apps that have positive estimated GPU savings except for the apps that
                are "Not Applicable"
        :param gpu_cluster_recommendation: The type of GPU cluster recommendation to generate.
               It accepts one of the following ("CLUSTER", "JOB" and the default value "MATCH").
                "MATCH": keep GPU cluster same number of nodes as CPU cluster;
                "CLUSTER": recommend optimal GPU cluster by cost for entire cluster;
                "JOB": recommend optimal GPU cluster by cost per job
        :param jvm_heap_size: The maximum heap size of the JVM in gigabytes
        :param verbose: True or False to enable verbosity to the wrapper script
        :param rapids_options: A list of valid Qualification tool options.
                Note that the wrapper ignores ["output-directory", "platform"] flags, and it does not support
                multiple "spark-property" arguments.
                For more details on Qualification tool options, please visit
                https://nvidia.github.io/spark-rapids/docs/spark-qualification-tool.html#qualification-tool-options
        """
        if verbose:
            # when debug is set to true set it in the environment.
            ToolLogging.enable_debug_mode()
        wrapper_qual_options = {
            'platformOpts': {
                'profile': profile,
                'deployMode': DeployMode.LOCAL,
            },
            'migrationClustersProps': {
                'cpuCluster': cpu_cluster,
                'gpuCluster': gpu_cluster
            },
            'jobSubmissionProps': {
                'remoteFolder': remote_folder,
                'platformArgs': {
                    'jvmMaxHeapSize': jvm_heap_size
                }
            },
            'eventlogs': eventlogs,
            'filterApps': filter_apps,
            'toolsJar': tools_jar,
            'gpuClusterRecommendation': gpu_cluster_recommendation
        }
        QualificationAsLocal(platform_type=CspEnv.EMR,
                             cluster=None,
                             output_folder=local_folder,
                             wrapper_options=wrapper_qual_options,
                             rapids_options=rapids_options).launch()

    @staticmethod
    def profiling(gpu_cluster: str = None,
                  worker_info: str = None,
                  eventlogs: str = None,
                  profile: str = None,
                  local_folder: str = None,
                  remote_folder: str = None,
                  tools_jar: str = None,
                  jvm_heap_size: int = 24,
                  verbose: bool = False,
                  **rapids_options) -> None:
        """
        The Profiling tool analyzes both CPU or GPU generated event logs and generates information
        which can be used for debugging and profiling Apache Spark applications.

        :param  gpu_cluster: The EMR-cluster on which the Spark applications were executed. The argument
                can be an EMR-cluster or a valid path to the cluster's properties file (json format)
                generated by the AWS CLI. If missing, then the argument worker_info has to be provided.
        :param  worker_info: A path pointing to a yaml file containing the system information of a
                worker node. It is assumed that all workers are homogenous.
                If missing, the wrapper pulls the worker info from the "gpu_cluster".
        :param  eventlogs: Event log filenames or S3 storage directories
                containing event logs (comma separated). If missing, the wrapper Reads the Spark's
                property `spark.eventLog.dir` defined in `gpu_cluster`. This property should be included
                in the output of `aws emr describe-cluster`. Note that the wrapper will raise an exception
                if the property is not set.
        :param profile: A named AWS profile to get the settings/credentials of the AWS account.
        :param local_folder: Local work-directory path to store the output and to be used as root
                directory for temporary folders/files. The final output will go into a subdirectory called
                ${local_folder}/prof-${EXEC_ID} where exec_id is an auto-generated unique identifier of the
                execution. If the argument is NONE, the default value is the env variable
                RAPIDS_USER_TOOLS_OUTPUT_DIRECTORY if any; or the current working directory.
        :param remote_folder: A S3 folder where the output is uploaded at the end of execution.
                If no value is provided, the output will be only available on local disk.
        :param tools_jar: Path to a bundled jar including Rapids tool. The path is a local filesystem,
                or remote S3 url. If missing, the wrapper downloads the latest rapids-4-spark-tools_*.jar
                from maven repo.
        :param verbose: True or False to enable verbosity to the wrapper script.
        :param jvm_heap_size: The maximum heap size of the JVM in gigabytes.
        :param rapids_options: A list of valid Profiling tool options.
                Note that the wrapper ignores ["output-directory", "worker-info"] flags, and it does not support
                multiple "spark-property" arguments.
                For more details on Profiling tool options, please visit
                https://nvidia.github.io/spark-rapids/docs/spark-profiling-tool.html#profiling-tool-options
        """
        if verbose:
            # when debug is set to true set it in the environment.
            ToolLogging.enable_debug_mode()
        wrapper_prof_options = {
            'platformOpts': {
                'profile': profile,
                'deployMode': DeployMode.LOCAL,
            },
            'migrationClustersProps': {
                'gpuCluster': gpu_cluster
            },
            'jobSubmissionProps': {
                'remoteFolder': remote_folder,
                'platformArgs': {
                    'jvmMaxHeapSize': jvm_heap_size
                }
            },
            'eventlogs': eventlogs,
            'toolsJar': tools_jar,
            'autoTunerFileInput': worker_info
        }
        ProfilingAsLocal(platform_type=CspEnv.EMR,
                         output_folder=local_folder,
                         wrapper_options=wrapper_prof_options,
                         rapids_options=rapids_options).launch()

    @staticmethod
    def bootstrap(cluster: str,
                  profile: str = None,
                  output_folder: str = None,
                  dry_run: bool = True,
                  key_pair_path: str = None,
                  verbose: bool = False) -> None:
        """
        Bootstrap tool analyzes the CPU and GPU configuration of the EMR cluster
        and updates the Spark default configuration on the cluster's master nodes.

        :param cluster: Name of the EMR cluster running an accelerated computing instance class g4dn.*
        :param profile: A named AWS profile to get the settings/credentials of the AWS account.
        :param output_folder: Local path where the final recommendations will be saved.
               Note that this argument only accepts local filesystem. If the argument is NONE,
               the default value is the env variable "RAPIDS_USER_TOOLS_OUTPUT_DIRECTORY" if any;
               or the current working directory.
        :param dry_run: True or False to update the Spark config settings on EMR master node.
        :param key_pair_path: A '.pem' file path that enables to connect to EC2 instances using SSH.
               If missing, the wrapper reads the env variable 'RAPIDS_USER_TOOLS_KEY_PAIR_PATH' if any.
               For more details on creating key pairs,
               visit https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/create-key-pairs.html.
        :param verbose: True or False to enable verbosity to the wrapper script.
        """
        if verbose:
            # when debug is set to true set it in the environment.
            ToolLogging.enable_debug_mode()
        wrapper_boot_options = {
            'platformOpts': {
                'profile': profile,
                'keyPairPath': key_pair_path
            },
            'dryRun': dry_run
        }
        bootstrap_tool = Bootstrap(platform_type=CspEnv.EMR,
                                   cluster=cluster,
                                   output_folder=output_folder,
                                   wrapper_options=wrapper_boot_options)
        bootstrap_tool.launch()

    @staticmethod
    def diagnostic(cluster: str,
                   profile: str = None,
                   output_folder: str = None,
                   key_pair_path: str = None,
                   thread_num: int = 3,
                   yes: bool = False,
                   verbose: bool = False) -> None:
        """
        Diagnostic tool to collect information from EMR cluster, such as OS version, # of worker nodes,
        Yarn configuration, Spark version and error logs etc. Please note, some sensitive information might
        be collected by this tool, e.g. access secret configured in configuration files or dumped to log files.
        :param cluster: Name of the EMR cluster running an accelerated computing instance class g4dn.*
        :param profile: A named AWS profile to get the settings/credentials of the AWS account.
        :param output_folder: Local path where the archived result will be saved.
               Note that this argument only accepts local filesystem. If the argument is NONE,
               the default value is the env variable "RAPIDS_USER_TOOLS_OUTPUT_DIRECTORY" if any;
               or the current working directory.
        :param key_pair_path: A '.pem' file path that enables to connect to EC2 instances using SSH.
               If missing, the wrapper reads the env variable 'RAPIDS_USER_TOOLS_KEY_PAIR_PATH' if any.
               For more details on creating key pairs,
               visit https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/create-key-pairs.html.
        :param thread_num: Number of threads to access remote cluster nodes in parallel. The valid value
               is 1~10. The default value is 3.
        :param yes: auto confirm to interactive question.
        :param verbose: True or False to enable verbosity to the wrapper script.
        """
        if verbose:
            # when debug is set to true set it in the environment.
            ToolLogging.enable_debug_mode()
        wrapper_diag_options = {
            'platformOpts': {
                'profile': profile,
                'keyPairPath': key_pair_path,
            },
            'threadNum': thread_num,
            'yes': yes,
        }
        diag_tool = Diagnostic(platform_type=CspEnv.EMR,
                               cluster=cluster,
                               output_folder=output_folder,
                               wrapper_options=wrapper_diag_options)
        diag_tool.launch()


class EMRWrapper:  # pylint: disable=too-few-public-methods
    """
    A wrapper script to run RAPIDS Accelerator tools (Qualification, Profiling, and Bootstrap) on Amazon EMR.
    """

    def __init__(self):
        self.qualification = CliEmrLocalMode.qualification
        self.profiling = CliEmrLocalMode.profiling
        self.bootstrap = CliEmrLocalMode.bootstrap
        self.diagnostic = CliEmrLocalMode.diagnostic
