# SPDX-FileCopyrightText: Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.

import contextlib
import importlib
import inspect
import json
import os
import psutil
import shutil
import signal
import sys

from datetime import datetime
from enum import Enum
from functools import partial

from nsys_recipe import nsys_constants
from nsys_recipe.lib import args, exceptions, nsys_path
from nsys_recipe.log import logger


class Mode(Enum):
    """Data processing modes."""

    NONE = 0
    CONCURRENT = 1
    DASK_FUTURES = 2


class Context:
    """Abstract base class for data processing engines, each associated to a Mode."""

    mode = None
    context_map = None

    def __init__(self):
        def exit_signal(signum, frame):
            sys.exit(128 + signum)

        signal.signal(signal.SIGINT, exit_signal)

    def __enter__(self):
        return self

    def __exit__(*args):
        pass

    def launch(self, function, *args, **kwargs):
        """Wrapper for task execution.

        Parameters
        ----------
        function : callable
            Function to execute.
        *args : tuple
            Positional arguments.
        **kwargs : dict
            Dictionary of keyword arguments.

        Returns
        -------
        result : return type of function
        """
        raise NotImplementedError

    def map(self, function, *iterables, **kwargs):
        """Execute task for each iterable.

        The function is applied to each element in iterables, which should
        have the same length. The kwargs remains packed and is passed as
        argument to the function. Results are guaranteed to be in the same
        order as the input.

        Parameters
        ----------
        function : callable
            Function to execute.
        *iterables : iterables
            Objects to map over.
        **kwargs : dict
            Dictionary of keyword arguments.

        Returns
        -------
        result : list
        """
        raise NotImplementedError

    def wait(self, waitable):
        """Wrapper for task completion.

        If waitable is a remote result, wait until computation completes and
        return the value of the variable. This is meant to be used on results
        of launch and map functions.

        Returns
        -------
        result : depends on input
            Returns a list of results if the waitable is a list.
            Returns single element otherwise.
        """
        raise NotImplementedError

    @staticmethod
    def import_module(name):
        try:
            module = importlib.import_module(name)
        except ModuleNotFoundError as e:
            raise exceptions.ModeModuleNotFoundError(e) from e
        return module

    @staticmethod
    def register_process_termination():
        def terminate_process(signum, frame):
            psutil.Process().terminate()

        signal.signal(signal.SIGINT, terminate_process)

    @staticmethod
    def convert_to_original_type(value):
        with contextlib.suppress(ValueError):
            return int(value)
        with contextlib.suppress(ValueError):
            return float(value)
        with contextlib.suppress(ValueError):
            return bool(value)

        return value

    @staticmethod
    def get_mode_argument_dict(prefix):
        argument_dict = {}

        for key, value in os.environ.items():
            if key.upper().startswith(prefix):
                # Lowercase and remove the prefix.
                processed_key = key.lower()[len(prefix) :]
                argument_dict[processed_key] = Context.convert_to_original_type(value)

        return argument_dict

    @classmethod
    def create_context(cls, mode=Mode.CONCURRENT):
        """Create an instance of Context corresponding to mode.

        The first time this is called, create context_map that maps each
        context to its mode.

        Parameters
        ----------
        mode : Mode
            Mode of the context to create.

        Returns
        -------
        context : Context
        """
        if cls.context_map is None:
            keys = Mode
            values = [ContextNone, ContextConcurrent, ContextDaskFutures]
            cls.context_map = dict(zip(keys, values))

        if mode not in cls.context_map:
            raise NotImplementedError

        return cls.context_map[mode]()


class ContextNone(Context):
    """Standard single-threaded mode."""

    mode = Mode.NONE

    def launch(self, function, *args, **kwargs):
        return function(*args, **kwargs)

    def map(self, function, *iterables, **kwargs):
        partial_func = partial(function, **kwargs)
        return [*map(partial(self.launch, partial_func), *iterables)]

    def wait(self, waitable):
        return waitable


class ContextConcurrent(Context):
    """Concurrent mode using concurrent.futures."""

    mode = Mode.CONCURRENT

    def __init__(self, executor=None):
        super().__init__()

        self._executor = executor
        self._custom = False

        if self._executor is None:
            self._custom = True
            argument_dict = self.get_mode_argument_dict("NSYS_CONCURRENT_")
            pkg_concurrent_futures = Context.import_module("concurrent.futures")

            # The 'initializer' argument is available starting from 3.7.
            if sys.version_info >= (3, 7):
                self._executor = pkg_concurrent_futures.ProcessPoolExecutor(
                    initializer=Context.register_process_termination, **argument_dict
                )
            else:
                self._executor = pkg_concurrent_futures.ProcessPoolExecutor(
                    **argument_dict
                )

    def __enter__(self):
        if self._executor is None:
            raise RuntimeError("Executor is shutdown.")
        return self

    def __exit__(self, *args):
        self.close()

    def close(self):
        if self._custom:
            self._executor.shutdown()
            self._executor = None

    def launch(self, function, *args, **kwargs):
        return self._executor.submit(function, *args, **kwargs).result()

    def map(self, function, *iterables, **kwargs):
        partial_func = partial(function, **kwargs)
        return [*self._executor.map(partial_func, *iterables)]

    def wait(self, waitable):
        return waitable


class ContextDaskFutures(Context):
    """Concurrent mode using dask.distributed."""

    mode = Mode.DASK_FUTURES

    def __init__(self, cluster=None):
        super().__init__()

        self._cluster = cluster

        pkg_dask_distributed = Context.import_module("distributed")

        if self._cluster is not None:
            self._client = pkg_dask_distributed.Client(self._cluster)
        else:
            argument_dict = self.get_mode_argument_dict("NSYS_DASK_")
            self._client = pkg_dask_distributed.Client(**argument_dict)

        def _dask_worker_callback(recipe_pkg_path):
            Context.register_process_termination()
            sys.path.insert(0, recipe_pkg_path)

        # We assume that all worker nodes have the same recipe path.
        recipe_pkg_path = nsys_path.find_installed_file(
            nsys_constants.NSYS_PYTHON_PKG_DIR
        )
        partial_callback = partial(_dask_worker_callback, recipe_pkg_path)
        self._client.register_worker_callbacks(setup=partial_callback)

    def __enter__(self):
        if self._client is None:
            raise RuntimeError("Client is closed.")
        return self

    def __exit__(self, *args):
        self.close()

    def close(self):
        if self._client is not None:
            self._client.close()
            self._client = None

    def launch(self, function, *args, **kwargs):
        return self._client.submit(function, *args, **kwargs)

    def map(self, function, *iterables, **kwargs):
        partial_func = partial(function, **kwargs)
        return self._client.map(partial_func, *iterables)

    def wait(self, waitable):
        if isinstance(waitable, list):
            return self._client.gather(waitable)
        return waitable.result()


class Recipe:
    """Base class for all recipes"""

    recipe_dir = None
    recipe_name = None
    metadata_dict = None

    def __init__(self, parsed_args):
        """Initialize.

        Parameters
        ----------
        parsed_args : argparse.Namespace
            Parsed arguments.
        """
        self._parsed_args = parsed_args
        self._output_dir = None
        self._output_files = {}
        self._analysis_dict = {}
        # Used as the default output directory if the user doesn't specify a
        # name using the '--output' option.
        self._default_output_name = self.get_recipe_name()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.output_dir is not None and exc_type is not None:
            shutil.rmtree(self._output_dir)

    @property
    def output_dir(self):
        """Get output directory name without creating it."""
        return self._output_dir

    def set_default_output_name(self, name):
        """Set the default output directory name."""
        self._default_output_name = name

    def _generate_incremental_name(self, name):
        i = 1
        while os.path.exists(f"{name}-{i}"):
            i += 1
        return f"{name}-{i}"

    def _generate_unique_output_name(self):
        output_name = self.get_parsed_arg("output")
        if output_name is None:
            return self._generate_incremental_name(self._default_output_name)

        if os.path.exists(output_name):
            if not self.get_parsed_arg("force_overwrite", False):
                logger.warning(
                    f"Failed to create '{output_name}': Directory exists."
                    " Use '--force-overwrite' to overwrite existing directories,"
                    " or '--auto-increment' for a unique name."
                )
                return self._generate_incremental_name(self._default_output_name)
            else:
                shutil.rmtree(output_name)

        return output_name

    def get_parsed_arg(self, name, default=None):
        return getattr(self._parsed_args, name, default)

    def get_output_dir(self):
        """Return a unique output directory name.

        The first time this is called, create a unique directory where the
        output files are stored.
        Unless the '--output' option is specified with a directory name that
        does not exist or used along the '--force-overwrite' option, a unique
        directory name will be generated with an incrementing id.
        """
        if self._output_dir is None:
            self._output_dir = self._generate_unique_output_name()
            os.makedirs(self._output_dir)

        return self._output_dir

    def add_output_file(self, filename, filetype=None):
        """Get path of the output file.

        Prepend the output directory name to filename.
        If filetype is not None, add it to _output_files so it can later be
        recorded in the nsys-analysis json file.

        Parameters
        ----------
        filename : str
            Output file name.
        filetype : str
            File type to be recorded.

        Returns
        -------
        filepath : str
            Output file path.
        """
        if filetype:
            self._output_files[filename] = filetype
        return os.path.join(self.get_output_dir(), filename)

    def create_analysis_file(self):
        """Create the nsys-analysis json file containing metadata."""
        output_name = os.path.basename(self.get_output_dir())
        analysis_filename = f"{output_name}.nsys-analysis"

        with open(self.add_output_file(analysis_filename), "w") as f:
            json.dump(self._analysis_dict, f, indent=4)

    def create_notebook(self, notebook_name, dir_path=None, replace_dict=None):
        """Create output jupyter notebook from an existing template notebook.

        The output notebook is created under the same name as the template.
        Any key strings contained in replace_dict will be replaced by its value.

        Parameters
        ----------
        notebook_name : str
            Name of the template notebook file located in the same directory as
            the recipe script.
        dir_path : str
            Directory path to the notebook in case it is located in a different
            location.
        replace_dict : dict
            Dictionary containing the string to be replaced and the new value.
        """

        if dir_path is not None:
            notebook_dir = dir_path
        else:
            notebook_dir = os.path.dirname(inspect.getmodule(self).__file__)

        nb_template = os.path.join(notebook_dir, notebook_name)
        nb_output_file = self.add_output_file(notebook_name, "Notebook")

        if not os.path.exists(nb_template):
            raise NotImplementedError(f"File {nb_template} not found.")

        if replace_dict:
            with open(nb_template, "r") as f:
                file_content = f.read()

            for key, value in replace_dict.items():
                file_content = file_content.replace(str(key), str(value))

            with open(nb_output_file, "w") as f:
                f.write(file_content)
        else:
            shutil.copy(nb_template, nb_output_file)

    def add_notebook_helper_file(self, filename):
        """Copy the helper file from the lib directory to the output directory,
        where it can be accessed by the notebook generated by the
        create_notebook function.

        Parameters
        ----------
        filename : str
            Name of the helper file located in the lib directory.
        """
        filepath = os.path.join(os.path.dirname(__file__), filename)
        if not os.path.exists(filepath):
            raise NotImplementedError(f"File {filepath} not found.")

        output_path = self.add_output_file(filename)
        shutil.copy(filepath, output_path)

    def run(self, context):
        self._analysis_dict = {
            "RecipeJsonVersion": "1.2.0",
            "RecipeRuntimeVersion": "2.0.0",
            "RecipeName": self.get_recipe_name(),
            "DisplayName": self.get_metadata("display_name"),
            "Mode": context.mode.name.replace("_", "-").lower(),
            "StartTime": str(datetime.now()),
        }

    @classmethod
    def get_recipe_dir(cls):
        if cls.recipe_dir is None:
            cls.recipe_dir = os.path.dirname(inspect.getmodule(cls).__file__)
        return cls.recipe_dir

    @classmethod
    def get_recipe_name(cls):
        if cls.recipe_name is None:
            cls.recipe_name = os.path.basename(cls.get_recipe_dir())
        return cls.recipe_name

    @classmethod
    def get_metadata(cls, key):
        """Get metadata info.

        Parameters
        ----------
        key : str
            Key to be searched in the metadata dictionary item.

        Returns
        -------
        value : str
            Value for the specified key if key is found, None otherwise.
        """
        if cls.metadata_dict is None:
            json_path = os.path.join(cls.get_recipe_dir(), "metadata.json")

            if not os.path.exists(json_path):
                raise NotImplementedError("File metadata.json not found.")

            with open(json_path) as f:
                cls.metadata_dict = json.load(f)

        if key not in cls.metadata_dict:
            raise NotImplementedError(f"Key {key} not found in metadata.json.")

        val = cls.metadata_dict[key]
        # A list of strings was used to save the description for better
        # readability. Convert the list into a single string before returning.
        if key == "description":
            return "".join(val)
        return val

    @classmethod
    def get_argument_parser(cls):
        """Get default argument parser."""
        return args.ArgumentParser(
            prog=cls.get_recipe_name(),
            description=cls.get_metadata("description"),
            formatter_class=args.TextHelpFormatter,
        )