# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#

import typing as tp
import warnings
from pathlib import Path
from typing import Any, List, Optional, Type, Union

from ..core import plugins
from ..core.core import Executor, Job
from ..core.utils import DelayedSubmission


def _convert_deprecated_args(kwargs: tp.Dict[str, Any], deprecated_args: tp.Mapping[str, str]) -> None:
    for arg in list(kwargs):
        new_arg = deprecated_args.get(arg)
        if not new_arg:
            continue
        kwargs[new_arg] = kwargs.pop(arg)
        warnings.warn(f"Setting '{arg}' is deprecated. Use '{new_arg}' instead.")


class AutoExecutor(Executor):
    """Automatic job executor
    This class is used to hold the parameters to run a job either on the cluster
    corresponding to the environment.
    It can also be used to run job locally or in debug mode.
    In practice, it will create a bash file in the specified directory for each job,
    and pickle the task function and parameters. At completion, the job will also pickle
    the output. Logs are also dumped in the same directory.

    Executor specific parameters must be specified by prefixing them with the name
    of the executor they refer to. eg:
        - 'chronos_conda_file' (internal)
        - 'slurm_max_num_timeout'
    See each executor documentation for the list of available parameters.

    Parameters
    ----------
    folder: Path/str
        folder for storing job submission/output and logs.
    warn_ignored: bool
        prints a warning each time a parameter is provided but ignored because it is only
        useful for the other cluster.
    cluster: str
        Forces AutoExecutor to use the given environment. Use "local" to run jobs locally,
        "debug" to run jobs in process.
    kwargs: other arguments must be prefixed by the name of the executor they refer to.
        {exname}_{argname}: see {argname} documentation in {Exname}Executor documentation.

    Note
    ----
    - be aware that the log/output folder will be full of logs and pickled objects very fast,
      it may need cleaning.
    - use update_parameters to specify custom parameters (gpus_per_node etc...). If you
      input erroneous parameters, an error will print all parameters available for you.
    """

    _ctor_deprecated_args = {"max_num_timeout": "slurm_max_num_timeout", "conda_file": "chronos_conda_file"}

    def __init__(self, folder: Union[str, Path], cluster: Optional[str] = None, **kwargs: Any) -> None:
        self.cluster = cluster or self.which()

        executors = plugins.get_executors()
        if self.cluster not in executors:
            raise ValueError(f"AutoExecutor doesn't know any executor named {self.cluster}")

        _convert_deprecated_args(kwargs, self._ctor_deprecated_args)
        err = "Extra arguments must be prefixed by executor named, received unknown arg"
        err_ex_list = f"Known executors: {', '.join(executors)}."
        for name in kwargs:
            assert "_" in name, f"{err} '{name}'. {err_ex_list}"
            prefix = name.split("_")[0]
            assert (
                prefix in executors
            ), f"{err} '{name}', and '{prefix}' executor is also unknown. {err_ex_list}"
        self._executor = flexible_init(executors[self.cluster], folder, **kwargs)

        valid = self._valid_parameters()
        self._deprecated_args = {
            arg: f"{ex_name}_{arg}"
            for ex_name, ex in executors.items()
            for arg in ex._valid_parameters()
            if arg not in valid
        }
        super().__init__(self._executor.folder, self._executor.parameters)

    @staticmethod
    def which() -> str:
        """Returns what is the detected cluster."""
        executors = plugins.get_executors()
        best_ex = max(executors, key=lambda ex: executors[ex].affinity())

        if executors[best_ex].affinity() <= 0:
            raise RuntimeError(f"Did not found an available executor among {executors.keys()}.")

        return best_ex

    def register_dev_folders(self, folders: List[Union[str, Path]]) -> None:
        """Archive a list of folders to be untarred in the job working directory.
        This is only implemented for internal cluster, for running job on non-installed packages.
        This is not useful on slurm since the working directory of jobs is identical to
        your work station working directory.

        folders: list of paths
            The list of folders to archive and untar in the job working directory
        """
        register = getattr(self._executor, "register_dev_folders", None)
        if register is not None:
            register(folders)
        else:
            # TODO this should be done through update parameters
            warnings.warn(
                "Ignoring dev folder registration as it is only supported (and needed) for internal cluster"
            )

    @classmethod
    def _valid_parameters(cls) -> tp.Set[str]:
        return {"name", "timeout_min", "mem_gb", "nodes", "cpus_per_task", "gpus_per_node", "tasks_per_node"}

    def _internal_update_parameters(self, **kwargs: Any) -> None:
        """Updates submission parameters to srun/crun.

        Parameters
        ----------
        AutoExecutors provides shared parameters that are translated for each specific cluster.
        Those are: timeout_min (int), mem_gb (int), gpus_per_node (int), cpus_per_task (int),
        nodes (int), tasks_per_node (int) and name (str).
        Cluster specific parameters can be specified by prefixing them with the cluster name.

        Notes
        -----
        - Cluster specific parameters win over shared parameters.
            eg: if both `slurm_time` and `timeout_min` are provided, then:
                - `slurm_time` is used on the slurm cluster
                - `timeout_min` is used on other clusters
        """
        # check type of replaced variables
        generics = AutoExecutor._valid_parameters()
        for name in generics:
            if name in kwargs:
                expected_type = int if name != "name" else str
                assert isinstance(kwargs[name], expected_type), (
                    f'Parameter "{name}" expected type {expected_type} ' f'(but value: "{kwargs[name]}")'
                )

        _convert_deprecated_args(kwargs, self._deprecated_args)
        specific = [x.split("_", 1) for x in kwargs if x not in generics]

        invalid = []
        executors = plugins.get_executors()
        for ex_arg in specific:
            if len(ex_arg) != 2:
                invalid.append(f"Parameter '{ex_arg[0]}' need to be prefixed by an executor name.")
                continue
            ex, arg = ex_arg

            if ex not in executors:
                invalid.append(f"Unknown executor '{ex}' in parameter '{ex}_{arg}'.")
                continue

            valid = executors[ex]._valid_parameters()
            if arg not in valid | generics:
                invalid.append(
                    f"Unknown argument '{arg}' for executor '{ex}' in parameter '{ex}_{arg}'."
                    + " Valid arguments: "
                    + ", ".join(valid)
                )
                continue
        if invalid:
            invalid.append(f"Known executors: {', '.join(executors.keys())}")
            raise NameError("\n".join(invalid))

        # add cluster specific generic overrides
        kwargs.update(
            **{
                arg: kwargs.pop(f"{ex}_{arg}")
                for ex, arg in specific
                if ex == self.cluster and arg in generics
            }
        )
        parameters = self._executor._convert_parameters({k: kwargs[k] for k in kwargs if k in generics})
        # update parameters in the core executor
        for (ex, arg) in specific:
            # update cluster specific non-generic arguments
            if arg not in generics and ex == self.cluster:
                parameters[arg] = kwargs[f"{ex}_{arg}"]

        self._executor._internal_update_parameters(**parameters)

    def _internal_process_submissions(
        self, delayed_submissions: tp.List[DelayedSubmission]
    ) -> tp.List[Job[tp.Any]]:
        return self._executor._internal_process_submissions(delayed_submissions)


def flexible_init(cls: Type[Executor], folder: Union[str, Path], **kwargs: Any) -> Executor:
    prefix = cls.name() + "_"
    return cls(folder, **{k[len(prefix) :]: kwargs[k] for k in kwargs if k.startswith(prefix)})
