import inspect
import os
import pickle
import shutil
import subprocess
import tempfile
import time
import sys
import cloudpickle
from .config import make_config
from .execution_contexts import DefaultContext
from .db import job_db_factory
from .util import slurm_available
import warnings


class SlurmSignatureMismatchException(Exception):
    pass


class Slurm:
    """Map a function to be executed with Slurm workload manager.
    Reads a config file (if it exists) in the home directory which should look as the default
    in parallel.config.

    The mapper reads commonly used parameters from a configuration file stord in ``~/.parallel``.
    An example configuration file could look as follows:


    .. code::

        #~/.parallel
        [DIRECTORIES]
        TMP=/tmp

        [BROKER]
        TYPE=REDIS   # can be SQLITE or REDIS

        [SLURM]
        PARTITION=p.gaba
        NICENESS=0

        [REDIS]
        HOST=127.0.0.1


    Parameters
    ----------

    tmp_directory: str or None
        Directory where temporary job pickle files are stored.
        If set to None a tmp directory is read from the ''~/.parallel'' configuration file.
        It this file does not exist a tmp directory within the user home directory is created.

    memory: str
        RAM requested by each job, e.g. "10G".

    time_h: int
        Job run time in hours.

    python_executable_path: str or None
        The python interpreter which executes the jobs.
        If set to None, the currently executing interpreter is used as returned by ``sys.executable``.

    slurm_error_file: str or None
        File to which stderr messages from workers are stored.
        If set to None, a file within the tmp_directory is used.

    slurm_output_file: str or None
        File to which stdout messages from workers are stored.
        If set to None, a file within the tmp_directory is used.

    name: str
        A name for the job.

    partition: str
        The Slurm partition.

    niceness: int
        Slurm job niceness.
        Default: 0

    num_threads: int, default = 1
        Number of threads for each worker.
        This also sets the environment variable MKL_NUM_THREADS, OMP_NUM_THREADS to the
        sepcified number to handle jobs which use OpenMP etc. correctly.

    execution_context: :class:`DefaultContext <parallel.execution_contexts.DefaultContext>`, :class:`ProfilingContext <parallel.execution_contexts.ProfilingContext>`, :class:`NamedPrinter <parallel.execution_contexts.NamedPrinter>`
        Any context manager can be passed here.
        The ``__enter__`` method is called before evaluating the function on the cluster.
        The ``__exit__`` method directly after the function run finished.

    chunk_size: int, default=1
        Number of tasks executed within one job.

            .. warning::

                If ``chunk_size`` is larger than 1, this can have bad side effects
                as all the jobs within one chunk are executed within the python
                process.


    Returns
    -------
    slurm: Slurm
        Configured Slurm mapper.
    """
    def __init__(self, tmp_directory: str =None, memory: str = '3G', time_h: int =100, name="map", 
                 python_executable_path: str =None, slurm_error_file: str =None, slurm_output_file:str =None,
                 partition=None, niceness=None, num_threads: int =1, execution_context=DefaultContext, chunk_size=1):

        # simple assignments
        self.memory = memory
        self.time_h = time_h

        self.config = make_config()

        if partition is not None:
            self.config["SLURM"]["PARTITION"] = partition
        if tmp_directory is not None:
            self.config["DIRECTORIES"]["TMP"] = tmp_directory
        if niceness is not None:
            self.config["SLURM"]["NICENESS"] = str(niceness)

        try:
            os.makedirs(self.config["DIRECTORIES"]["TMP"])
        except FileExistsError:
            pass

        self.time = str(time_h) + ":00:00"

        self.job_name = name
        self.num_threads = num_threads
        self.execution_context = execution_context
        self.chunk_size = chunk_size

        if chunk_size != 1:
            warnings.warn("Chunk size != 1. This can potentially have bad side effect.")

        if not slurm_available():
            print("Warning: Could not find Slurm installation.", file=sys.stderr)


        # python interpreter which executes the jobs
        if type(time_h) is not int:
            raise Exception('Time should be an integer hour')
        if python_executable_path is None:
            python_executable_path = sys.executable
        self.python_executable_path = python_executable_path

        # slurm stderr
        if slurm_error_file is not None:
            self.slurm_error_file = slurm_error_file
        else:
            self.slurm_error_file = "slurm_errors.txt"

        # slurm stdout
        if slurm_output_file is not None:
            self.slurm_output_file = slurm_output_file
        else:
            self.slurm_output_file = "slurm_output.txt"

    def __repr__(self):
        return ("<Slurm memory={} time={} niceness={} num_threads={} chunk_size={} tmp_dir={} python_executable={}>"
                .format(self.memory, self.time, self.config["SLURM"]["NICENESS"], self.num_threads, self.chunk_size,
                        self.config["DIRECTORIES"]["TMP"], self.python_executable_path))

    @staticmethod
    def _validate_function_arguments(function, array):
        """

        Parameters
        ----------
        function: callable
            Function to be mapped.
        array: list
            Array of values.

        Raises an Exception if arguments do not match.
        This prevents unnecessary Job submission and raises an error earlier.

        """
        signature = inspect.signature(function)
        try:
            for argument in array:
                signature.bind(argument)
        except TypeError as err:
            raise SlurmSignatureMismatchException("Your jobs were not submitted as the function could not be applied to the arguments.") from err

    def map(self, function, array):
        """
        Does what map(function, array) would do, but does it
        via an array job on Slurm by pickling objects, storing
        them in a temporary folder, submitting them to the Slurm and
        then reading and returning the results.

        Parameters
        ----------

        function: callable
            The function to be mapped.

        array: iterable
            The values to which the function is applied

        Returns
        -------

        result_list: list
            List of results of function application.
            This list can also contain ``Exception`` objects.
        """
        array = list(array)

        self._validate_function_arguments(function, array)
        tmp_dir = tempfile.mkdtemp(prefix="", suffix='_slurm_job', dir=self.config["DIRECTORIES"]["TMP"])

        # jobs directory
        jobs_dir = os.path.join(tmp_dir, 'jobs')
        os.mkdir(jobs_dir)

        # create results directory
        os.mkdir(os.path.join(tmp_dir, 'results'))

        # store the function
        with open(os.path.join(tmp_dir, 'function.pickle'), 'wb') as my_file:
            cloudpickle.dump(function, my_file)

        # store execution context
        with open(os.path.join(tmp_dir, 'ExecutionContext.pickle'), 'wb') as my_file:
            cloudpickle.dump(self.execution_context, my_file)

        # store the array
        for task_nr, start_index in enumerate(range(0, len(array), self.chunk_size)):
            with open(os.path.join(jobs_dir, str(task_nr+1) + '.job'), 'wb') as my_file:
                cloudpickle.dump(list(array[start_index:start_index+self.chunk_size]), my_file)

        nr_tasks = task_nr + 1

        batch_file = self._render_batch_file(nr_tasks, tmp_dir)

        with open(os.path.join(tmp_dir, 'job.sh'), 'w') as my_file:
            my_file.write(batch_file)


        # crate job jd
        job_db = job_db_factory(tmp_dir)
        job_db.create(len(array))

        # start the job with qsub
        subprocess.run(['sbatch', os.path.join(tmp_dir, 'job.sh')], stdout=subprocess.PIPE)

        # wait for the tasks to be finished
        finished = False
        while not finished:
            time.sleep(5)
            for k in range(nr_tasks)[::-1]:  # check from last to first job
                if job_db.wait_for_job(k+1, self.time_h):  # +1 offset for k b/c Slurm array jobs start at 1, not at 0
                    break
            else:
                finished = True
        # from here on all tasks are finished

        # sleep to make sure files are entirely written
        time.sleep(5)

        # make the results array
        results = []
        had_exception = False
        for task_nr in range(nr_tasks):
            try:
                my_file = open(os.path.join(tmp_dir, 'results', str(task_nr+1) + '.result'), 'rb')
                single_result = pickle.load(my_file)
                results += single_result
                my_file.close()
            except Exception as e:
                results.append(Exception('Could not load temporary result file:' + str(e)))
                had_exception = True

        # delete the temporary folder if there was no problem
        # and execution context does not need it
        if self.execution_context.keep_output_directory:
            pass
        elif had_exception:
            tmp_dir = tmp_dir[:-1] if tmp_dir[-1] == '/' else tmp_dir
            os.rename(tmp_dir, tmp_dir + '_with_exception')
        else:
            shutil.rmtree(tmp_dir)
            job_db.clean_up()
        return results

    def _render_batch_file(self, nr_tasks, tmp_dir):
        # create the file to be submitted to Slurm via sbatch
        # for array jobs the ressource request are per task!
        try:
            pythonpath = os.environ['PYTHONPATH']
        except KeyError:
            pythonpath = ""

        batch_file = """#!/bin/bash
#SBATCH --job-name="{job_name}-{map_name}"
#SBATCH --partition="{partition}"
#SBATCH --cpus-per-task={num_threads}
#SBATCH --mem="{RAM}"
#SBATCH --time="{time}"
#SBATCH --nice="{niceness}"
#SBATCH --export=ALL
#SBATCH --array=1-{nr_elements}
#SBATCH --error="{slurm_error_file}"
#SBATCH --output="{slurm_output_file}"
#SBATCH --chdir="{current_working_directory}"
export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
export MKL_NUM_THREADS=$SLURM_CPUS_PER_TASK
export PYTHONPATH="{pythonpath}"
{python_executable_path} -m parallel.execute_array_job {tmp_dir} $SLURM_ARRAY_TASK_ID
{python_executable_path} -m parallel.cleanup_array_job {tmp_dir} $SLURM_ARRAY_TASK_ID
""".format(tmp_dir=tmp_dir, nr_elements=nr_tasks,
           current_working_directory=os.getcwd(),
           RAM=self.memory, time=self.time,
           job_name=self.job_name,
           python_executable_path=self.python_executable_path,
           slurm_error_file=os.path.join(tmp_dir, self.slurm_error_file),
           slurm_output_file=os.path.join(tmp_dir, self.slurm_output_file),
           niceness=self.config["SLURM"]["NICENESS"],
           partition=self.config["SLURM"]["partition"],
           map_name=os.path.split(tmp_dir)[-1],
           num_threads=self.num_threads,
           pythonpath=pythonpath)
        return batch_file
