# 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
from pathlib import Path
from typing import IO, Any, Dict, List, Optional, Sequence, Union

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"}

LOCAL_REQUEUE_RETURN_CODE = 144


class LocalJob(core.Job[R]):
    def __init__(
        self,
        folder: Union[Path, str],
        job_id: str,
        tasks: Sequence[int] = (0,),
        process: Optional["subprocess.Popen['bytes']"] = None,
    ) -> None:
        super().__init__(folder, job_id, tasks)
        self._cancel_at_deletion = False
        self._process = process
        # downcast sub-jobs to get proper typing
        self._sub_jobs: Sequence["LocalJob[R]"] = self._sub_jobs
        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
        """
        assert self._process is not None
        return self._process.poll() is not None

    @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:
            return "UNKNOWN"

    def get_info(self) -> Dict[str, str]:
        """Returns information about the job as a dict.
        """
        assert self._process is not None
        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
        assert self._process is not None
        self._process.send_signal(signal.SIGINT)

    def _interrupt(self, timeout: bool = False) -> None:
        """Sends preemption or timeout signal to the job (for testing purpose)

        Parameter
        ---------
        timedout: bool
            Whether to trigger a job time-out (if False, it triggers preemption)
        """
        assert self._process is not None
        if not timeout:
            self._process.send_signal(signal.SIGTERM)
            time.sleep(0.001)
        self._process.send_signal(signal.SIGUSR1)

    def trigger_timeout(self) -> None:
        assert self._process is not None
        self._process.send_signal(signal.SIGUSR1)

    def __del__(self) -> None:
        if self._cancel_at_deletion:
            if not self.get_info().get("jobState") == "FINISHED":
                self.cancel(check=False)


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: Union[str, Path], max_num_timeout: int = 3) -> None:
        super().__init__(folder, max_num_timeout=max_num_timeout)
        # preliminary check
        indep_folder = utils.JobPaths.get_first_id_independent_folder(self.folder)
        indep_folder.mkdir(parents=True, exist_ok=True)

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

        Valid parameters are:
        - timeout_min (float)
        - gpus_per_node (int)
        - tasks_per_node (int)
        - nodes (int). Must be 1 if specified
        - signal_delay_s (int): USR1 signal delay before timeout

        Other parameters are ignored
        """
        if kwargs.get("nodes", 0) > 1:
            raise ValueError("LocalExecutor can use only one node. Use nodes=1")
        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)
        process = start_controller(
            folder=self.folder,
            command=command,
            tasks_per_node=ntasks,
            cuda_devices=",".join(str(k) for k in range(self.parameters.get("gpus_per_node", 0))),
            timeout_min=self.parameters.get("timeout_min", 2.0),
            signal_delay_s=self.parameters.get("signal_delay_s", 30),
        )
        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 f"{sys.executable} -u -m submitit.core._submit '{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: Union[bytes, str]) -> str:
        # Not used, but need an implementation
        return "0"

    def _make_submission_command(self, submission_file_path: Path) -> 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,
) -> "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_EXECUTOR="local",
        CUDA_AVAILABLE_DEVICES=cuda_devices,
    )
    process = subprocess.Popen(
        [sys.executable, "-m", "submitit.local._local", str(folder)], shell=False, env=env
    )
    return process


class Controller:
    """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.tasks: List[subprocess.Popen] = []  # type: ignore
        self.stdouts: List[IO[Any]] = []
        self.stderrs: List[IO[Any]] = []
        self.pid = str(os.getpid())
        self.folder = Path(folder)
        signal.signal(signal.SIGUSR1, self._forward_signal)
        signal.signal(signal.SIGTERM, self._forward_signal)

    def _forward_signal(self, signum: signal.Signals, *args: Any) -> None:  # pylint:disable=unused-argument
        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 = [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(
                    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) -> Sequence[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(signal.SIGUSR1)

            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()
