# 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 os
import shlex
import signal
import subprocess
import sys
import time
import typing as tp
from pathlib import Path

from ..core import core, job_environment, logger, utils
from ..core.core import R

# pylint: disable-msg=too-many-arguments
# VALID_KEYS = {"timeout_min", "gpus_per_node", "tasks_per_node", "signal_delay_s", "visible_gpus", "setup"}

LOCAL_REQUEUE_RETURN_CODE = 144

# global variable storing unfinished processes of pickled jobs
# in case we need to reload them later
_PROCESSES: tp.Dict[str, "subprocess.Popen['bytes']"] = {}


class LocalJob(core.Job[R]):
    def __init__(
        self,
        folder: tp.Union[Path, str],
        job_id: str,
        tasks: tp.Sequence[int] = (0,),
        process: tp.Optional["subprocess.Popen['bytes']"] = None,
    ) -> None:
        super().__init__(folder, job_id, tasks)
        self._cancel_at_deletion = False
        # downcast sub-jobs to get proper typing
        self._sub_jobs: tp.Sequence["LocalJob[R]"] = self._sub_jobs
        # set process (to self and subjobs)
        self._process = process
        for sjob in self._sub_jobs:
            sjob._process = process

    def done(self, force_check: bool = False) -> bool:  # pylint: disable=unused-argument
        """Override to avoid using the watcher"""
        state = self.get_info()["jobState"]
        return state != "RUNNING"

    @property
    def state(self) -> str:
        """State of the job"""
        try:
            return self.get_info().get("jobState", "unknown")
        # I don't what is the exception returned and it's hard to reproduce
        except Exception:  # pylint: disable=broad-except
            return "UNKNOWN"

    def get_info(self, mode: str = "force") -> tp.Dict[str, str]:  # pylint: disable=unused-argument
        """Returns information about the job as a dict."""
        if self._process is None:
            state = "NO PROCESS AND NO RESULT"
            if self.paths.result_pickle.exists():
                state = "FINISHED"
            return {"jobState": state}
        poll = self._process.poll()
        if poll is None:
            state = "RUNNING"
        elif poll < 0:
            state = "INTERRUPTED"
        else:
            state = "FINISHED"
        return {"jobState": state}

    def cancel(self, check: bool = True) -> None:  # pylint: disable=unused-argument
        if self._process is not None:
            self._process.send_signal(signal.SIGINT)

    def _interrupt(self) -> None:
        """Sends preemption / timeout signal to the job (for testing purpose)"""
        if self._process is not None:
            self._process.send_signal(LocalJobEnvironment._usr_sig())

    def __del__(self) -> None:
        if self._cancel_at_deletion:
            if not self.get_info().get("jobState") == "FINISHED":
                self.cancel(check=False)
        # let's clear the process dict if we know it's finished
        if self.paths.result_pickle.exists():
            _PROCESSES.pop(self.job_id, None)

    # # # # # pickling below # # # # #

    def __getstate__(self) -> tp.Any:
        out = dict(self.__dict__)
        out["_process"] = None
        if self._process is not None:
            _PROCESSES[self.job_id] = self._process
        return out

    def __setstate__(self, state: tp.Any) -> None:
        # Restore instance attributes
        self.__dict__.update(state)
        # recover process if it still exists
        self._process = _PROCESSES.get(self.job_id, None)


class LocalJobEnvironment(job_environment.JobEnvironment):
    _env = {
        "job_id": "SUBMITIT_LOCAL_JOB_ID",
        "num_tasks": "SUBMITIT_LOCAL_NTASKS",
        "num_nodes": "SUBMITIT_LOCAL_JOB_NUM_NODES",
        "node": "SUBMITIT_LOCAL_NODEID",
        "global_rank": "SUBMITIT_LOCAL_GLOBALID",
        "local_rank": "SUBMITIT_LOCAL_LOCALID",
    }

    def _requeue(self, countdown: int) -> None:
        jid = self.job_id
        logger.get_logger().info(f"Requeued job {jid} ({countdown} remaining timeouts)")
        sys.exit(LOCAL_REQUEUE_RETURN_CODE)  # should help noticing if need requeuing


class LocalExecutor(core.PicklingExecutor):
    """Local job executor
    This class is used to hold the parameters to run a job locally.
    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.

    The submission file spawn several processes (one per task), with a timeout.


    Parameters
    ----------
    folder: Path/str
        folder for storing job submission/output and logs.

    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 (n_gpus etc...).
    """

    job_class = LocalJob

    def __init__(
        self,
        folder: tp.Union[str, Path],
        max_num_timeout: int = 3,
        max_pickle_size_gb: float = 1.0,
        python: tp.Optional[str] = None,
    ) -> None:
        super().__init__(
            folder,
            max_pickle_size_gb=max_pickle_size_gb,
            max_num_timeout=max_num_timeout,
        )
        self.python = shlex.quote(sys.executable) if python is None else python
        # preliminary check
        indep_folder = utils.JobPaths.get_first_id_independent_folder(self.folder)
        indep_folder.mkdir(parents=True, exist_ok=True)

    @classmethod
    def _valid_parameters(cls) -> tp.Set[str]:
        """Parameters that can be set through update_parameters"""
        return {"setup"}

    def _internal_update_parameters(self, **kwargs: tp.Any) -> None:
        """Update the parameters of the Executor.

        Valid parameters are:
        - timeout_min (float)
        - gpus_per_node (int)
        - visible_gpus (Sequence[int])
        - tasks_per_node (int)
        - nodes (int). Must be 1 if specified
        - signal_delay_s (int): signal (lately: USR2) delay before timeout

        Other parameters are ignored
        """
        if kwargs.get("nodes", 0) > 1:
            raise ValueError("LocalExecutor can use only one node. Use nodes=1")
        gpus_requested = kwargs.get("gpus_per_node", 0)
        visible_gpus = kwargs.get("visible_gpus", ())
        if not isinstance(visible_gpus, tp.Sequence):
            raise ValueError(f"Provided visible_gpus={visible_gpus} is not an instance of Sequence.")
        if not all(isinstance(x, int) for x in visible_gpus):
            raise ValueError(f"Provided visible_gpus={visible_gpus} contains an element that is not an int.")
        if len(visible_gpus) > 0 and gpus_requested > len(visible_gpus):
            raise ValueError(
                f"{gpus_requested} gpus requested, but only {visible_gpus} were specified visible."
            )
        super()._internal_update_parameters(**kwargs)

    def _submit_command(self, command: str) -> LocalJob[R]:
        # Override this, because the implementation is simpler than for clusters like Slurm
        # Only one node is supported for local executor.
        ntasks = self.parameters.get("tasks_per_node", 1)
        n_gpus = self.parameters.get("gpus_per_node", 0)
        visible_gpus = self.parameters.get("visible_gpus", ())
        gpus = range(n_gpus) if visible_gpus == () else visible_gpus[:n_gpus]
        process = start_controller(
            folder=self.folder,
            command=command,
            tasks_per_node=ntasks,
            cuda_devices=",".join(str(k) for k in gpus),
            timeout_min=self.parameters.get("timeout_min", 2.0),
            signal_delay_s=self.parameters.get("signal_delay_s", 30),
            stderr_to_stdout=self.parameters.get("stderr_to_stdout", False),
            setup=self.parameters.get("setup", ()),
        )
        job: LocalJob[R] = LocalJob(
            folder=self.folder, job_id=str(process.pid), process=process, tasks=list(range(ntasks))
        )
        return job

    @property
    def _submitit_command_str(self) -> str:
        return " ".join([self.python, "-u -m submitit.core._submit", shlex.quote(str(self.folder))])

    def _num_tasks(self) -> int:
        nodes: int = 1
        tasks_per_node: int = self.parameters.get("tasks_per_node", 1)
        return nodes * tasks_per_node

    def _make_submission_file_text(self, command: str, uid: str) -> str:
        return ""

    @staticmethod
    def _get_job_id_from_submission_command(string: tp.Union[bytes, str]) -> str:
        # Not used, but need an implementation
        return "0"

    def _make_submission_command(self, submission_file_path: Path) -> tp.List[str]:
        # Not used, but need an implementation
        return []


def start_controller(
    folder: Path,
    command: str,
    tasks_per_node: int = 1,
    cuda_devices: str = "",
    timeout_min: float = 5.0,
    signal_delay_s: int = 30,
    stderr_to_stdout: bool = False,
    setup: tp.Sequence[str] = (),
) -> "subprocess.Popen['bytes']":
    """Starts a job controller, which is expected to survive the end of the python session."""
    env = dict(os.environ)
    env.update(
        SUBMITIT_LOCAL_NTASKS=str(tasks_per_node),
        SUBMITIT_LOCAL_COMMAND=command,
        SUBMITIT_LOCAL_TIMEOUT_S=str(int(60 * timeout_min)),
        SUBMITIT_LOCAL_SIGNAL_DELAY_S=str(int(signal_delay_s)),
        SUBMITIT_LOCAL_NODEID="0",
        SUBMITIT_LOCAL_JOB_NUM_NODES="1",
        SUBMITIT_STDERR_TO_STDOUT="1" if stderr_to_stdout else "",
        SUBMITIT_EXECUTOR="local",
        CUDA_VISIBLE_DEVICES=cuda_devices,
        SUBMITIT_LOCAL_WITH_SHELL="1" if setup else "",
    )
    # The LocalJob will be responsible to polling and ending this process.
    # pylint: disable=consider-using-with
    proc_cmd: tp.Any = [sys.executable, "-m", "submitit.local._local", str(folder)]
    need_shell = bool(setup)
    if need_shell:
        proc_cmd = " && ".join(list(setup) + [shlex.join(proc_cmd)])
    process = subprocess.Popen(proc_cmd, shell=need_shell, env=env)
    return process


class Controller:  # pragma: no cover
    """This controls a job:
    - instantiate each of the tasks
    - sends timeout signal
    - stops all tasks if one of them finishes
    - cleans up the tasks/closes log files when deleted
    """

    # pylint: disable=too-many-instance-attributes

    def __init__(self, folder: Path):
        self.ntasks = int(os.environ["SUBMITIT_LOCAL_NTASKS"])
        self.command = shlex.split(os.environ["SUBMITIT_LOCAL_COMMAND"])
        self.timeout_s = int(os.environ["SUBMITIT_LOCAL_TIMEOUT_S"])
        self.signal_delay_s = int(os.environ["SUBMITIT_LOCAL_SIGNAL_DELAY_S"])
        self.stderr_to_stdout = bool(os.environ["SUBMITIT_STDERR_TO_STDOUT"])
        self.tasks: tp.List[subprocess.Popen] = []  # type: ignore
        self.stdouts: tp.List[tp.IO[tp.Any]] = []
        self.stderrs: tp.List[tp.IO[tp.Any]] = []
        with_shell = bool(os.environ["SUBMITIT_LOCAL_WITH_SHELL"])
        self.pid = str(os.getppid() if with_shell else os.getpid())
        self.folder = Path(folder)
        signal.signal(signal.SIGTERM, self._forward_signal)  # type: ignore

    # pylint:disable=unused-argument
    def _forward_signal(self, signum: signal.Signals, *args: tp.Any) -> None:
        for task in self.tasks:
            try:
                task.send_signal(signum)  # sending kill signal to make sure everything finishes
            except Exception:
                pass

    def start_tasks(self) -> None:
        self.folder.mkdir(exist_ok=True)
        paths = [utils.JobPaths(self.folder, self.pid, k) for k in range(self.ntasks)]
        self.stdouts = [p.stdout.open("a") for p in paths]
        self.stderrs = self.stdouts if self.stderr_to_stdout else [p.stderr.open("a") for p in paths]
        for k in range(self.ntasks):
            env = dict(os.environ)
            env.update(
                SUBMITIT_LOCAL_LOCALID=str(k), SUBMITIT_LOCAL_GLOBALID=str(k), SUBMITIT_LOCAL_JOB_ID=self.pid
            )
            self.tasks.append(
                subprocess.Popen(  # pylint: disable=consider-using-with
                    self.command,
                    shell=False,
                    env=env,
                    stderr=self.stderrs[k],
                    stdout=self.stdouts[k],
                    encoding="utf-8",
                )
            )

    def kill_tasks(self) -> None:
        # try and be progressive in deletion...
        for sig in [signal.SIGINT, signal.SIGKILL]:
            self._forward_signal(sig)
            # if one is still alive after sigterm and sigint, try sigkill after 1s
            if sig == signal.SIGINT and any(t.poll() is None for t in self.tasks):
                time.sleep(0.001)
                if any(t.poll() is None for t in self.tasks):
                    time.sleep(1.0)  # wait a bit more
        self.tasks = []
        files = self.stdouts + self.stderrs
        self.stdouts, self.stderrs = [], []  # remove all instance references
        for f in files:
            f.close()

    def wait(self, freq: int = 24) -> tp.Sequence[tp.Optional[int]]:
        """Waits for all tasks to finish or to time-out.

        Returns
        -------
        Sequence[Optional[int]]:
            Exit codes of each task.
            Some tasks might still have not exited, but they will have received the "timed-out" signal.
        """
        assert self.tasks, "Nothing to do!"
        timeout = freq * self.timeout_s
        almost_timeout = freq * (self.timeout_s - self.signal_delay_s)

        # safer to keep a for loop :)
        for step in range(timeout):
            exit_codes = [t.poll() for t in self.tasks]
            if all(e is not None for e in exit_codes):
                return exit_codes

            if step == almost_timeout:
                self._forward_signal(LocalJobEnvironment._usr_sig())

            time.sleep(1.0 / freq)
        return [t.poll() for t in self.tasks]

    def run(self, max_retry: int = 6) -> None:
        # max_retry is a safety measure, the submission also have a timeout_countdown,
        # and will fail if it times out too many times.
        for _ in range(max_retry):
            try:
                self.start_tasks()
                exit_codes = self.wait()
                requeue = any(e == LOCAL_REQUEUE_RETURN_CODE for e in exit_codes)
                if not requeue:
                    break
            finally:
                self.kill_tasks()
