"""
This module contains executor-classes for pipelines generated by
a instance of :class:`doepipeline.generator.PipelineGenerator`.

Classes:
* :class:`BasePipelineExecutor` - Base-class for all executors. Implements
  basic interface and method to run pipeline collections.

Exceptions:
* :class:`CommandError`
* :class:`PipelineRunFailed`
"""
import abc
import logging
import time
import platform
import locale
from io import StringIO
from collections import OrderedDict
import pandas as pd

from doepipeline import utils



class CommandError(Exception):
    """
    Raised whenever a command fails.
    """


class PipelineRunFailed(Exception):
    """
    Raised when running pipeline failed.
    """


class BasePipelineExecutor(object):
    """ Base-class for pipeline executor classes.

    Implements parallel pipeline execution given pipeline collection
    from :class:`doepipeline.generator.PipelineGenerator`. Pipeline
    can either be run in parallel using Linux screens or batch-sripts.

    Attributes:
    :ivar workdir:
    :ivar base_script:
    :ivar base_log:
    :ivar run_in_batch:
    :ivar poll_interval:
    :ivar running_jobs:

    Class attributes:
    :cvar JOB_FINISHED:
    :cvar JOB_RUNNING:
    :cvar JOB_FAILED:
    """

    __metaclass__ = abc.ABCMeta

    JOB_FINISHED = 'job_finished'
    JOB_RUNNING = 'job_running'
    JOB_FAILED = 'job_failed'

    def __init__(self, workdir=None, poll_interval=10,
                 base_command=None, base_log=None, recovery_mode=False):
        assert workdir is None or isinstance(workdir, str) and workdir.strip(),\
            'path must be None or string'
        assert not isinstance(poll_interval, bool) and\
            isinstance(poll_interval, int) and poll_interval > 0,\
            'poll_interval must be positive integer'

        if base_command is not None:
            try:
                self.base_command = utils.validate_command(base_command)
            except AssertionError as e:
                raise ValueError('invalid base-command: ' + str(e))
        else:
             self.base_command = 'nohup {script} > {logfile} 2>&1'

        if base_log is not None:
            try:
                self.base_log = utils.validate_log_file(base_log)
            except AssertionError as e:
                raise ValueError('invalid base-log: ' + str(e))
        else:
            self.base_log = '{name}_step_{i}.log'

        self.is_setting_up = True
        self.recovery = recovery_mode
        self.workdir = workdir if workdir is not None else '.'
        self.poll_interval = poll_interval
        self.running_jobs = dict()
        self.has_workdir = False
        self.has_experiment_dirs = False
        self.encoding = locale.getpreferredencoding()

    @abc.abstractmethod
    def execute_command(self, command, watch=False, wait=False, **kwargs):
        """ Extend to execute the given command in current execution
        environment and returns echoed stdin, stdout, stderr.

        Base-class provides basic input validation and if `watch` is
        True, the current command is added the instances `current_jobs`.

        :param str command: Command to execute.
        :param bool watch: If True, add command to process to watch.
        :param kwargs: Additional keyword arguments.
        :returns: stdin, stdout, stderr
        :rtype: tuple[str]
        """
        try:
            assert isinstance(command, str) and command.strip(),\
                'command must be a string'
            assert isinstance(watch, bool), 'watch must be boolean'
            assert isinstance(wait, bool), 'wait must be boolean'
            if watch:
                assert 'job_name' in kwargs,\
                    'if watch is True, job_name must be given'
        except AssertionError as e:
            raise ValueError(str(e))

        logging.debug(('Execute command: "{}", watch={}, '
                       'wait={}, kwargs={}').format(command, watch,
                                                    wait, kwargs))

    @abc.abstractmethod
    def poll_jobs(self):
        """ Abstract method.

        Implement method to check status of running jobs.
        The method should return a tuple of two consisting
        of job-status and a message.

        Allowed statuses are:

        * `JOB_RUNNING` / "job_running"
        * `JOB_FINISHED` / "job_finished"
        * `JOB_FAILED` / "job_failed"

        :returns: status, message
        :rtype: str, str
        """

    def run_pipeline_collection(self, pipeline_collection):
        """

        :param pipeline_collection:
        :return: Pipeline results in a data-frame.
        :rtype: pandas.DataFrame
        """
        # Initialization..
        experiment_index = list()
        job_steps = OrderedDict((name, list()) for\
                                name in pipeline_collection['JOBNAMES'])
        env_variables = pipeline_collection['ENV_VARIABLES']
        setup = pipeline_collection['SETUP_SCRIPTS']
        reserved = ['ENV_VARIABLES', 'SETUP_SCRIPTS', 'RESULTS_FILE',
                    'WORKDIR', 'SLURM', 'JOBNAMES']
        self.workdir = pipeline_collection['WORKDIR']
        kwargs = {
            key.lower(): pipeline_collection[key] for key in reserved \
            if key in pipeline_collection \
            and key not in ('ENV_VARIABLES', 'SETUP_SCRIPTS',
                            'RESULTS_FILE', 'WORKDIR', 'JOBNAMES')
        }

        _items = pipeline_collection.items()
        pipeline_scripts = OrderedDict([(key, value) for key, value in _items\
                                        if key not in reserved])

        # Move to working-directory
        try:
            self.change_dir(self.workdir)
        except CommandError:
            logging.info('{} not found, creating directory'.format(self.workdir))
            self.make_dir(self.workdir)
            self.change_dir(self.workdir)

        self.has_workdir = True

        # Run setup-scripts in work-dir.
        if setup is not None:
            for script in setup:
                self.execute_command(script)

        logging.info('Creating job directories.')
        for job_name, scripts in pipeline_scripts.items():
            if job_name in reserved:
                # Don't treat special names as jobs.
                continue

            experiment_index.append(job_name)
            logging.debug('Creating directory: {}'.format(job_name))
            self.make_dir(str(job_name))

            for job_name, script in zip(job_steps, scripts):
                job_steps[job_name].append(script)

        self.has_experiment_dirs = True
        logging.info('Executing pipeline.')
        self.run_jobs(job_steps, experiment_index, env_variables, **kwargs)

        # Step into each work folder and collect pipeline results.
        return self._parse_results_file(experiment_index, pipeline_collection)

    @abc.abstractmethod
    def run_jobs(self, job_steps, experiment_index, env_variables, **kwargs):
        """ Abstract method.

        :param list[list] job_steps:
        :param list[str] experiment_index:
        :param env_variables:
        :return:
        """

    @abc.abstractmethod
    def read_file_contents(self, file_name, directory=None, **kwargs):
        """ Abstract method.

        Override to specify how to read file contents.

        :param str file_name: Path of file.
        :param str directory: Directory containing `file_name`
        :return: File-contents
        :rtype: str
        """

    @abc.abstractmethod
    def set_env_variables(self, env_variables):
        """ Abstract method.

        Override to specify how to set environment variables.

        :param dict env_variables: key-value pairs of env-variables to set.
        """

    def wait_until_current_jobs_are_finished(self):
        # Monitor job status.
        while 'running':
            status, msg = self.poll_jobs()
            if status == BasePipelineExecutor.JOB_FINISHED:
                self.running_jobs = dict()
                break
            elif status == BasePipelineExecutor.JOB_RUNNING:
                time.sleep(self.poll_interval)
            else:
                self.running_jobs = dict()
                logging.critical('Pipeline failed: "{}"'.format(msg))
                raise PipelineRunFailed(msg)

    def change_dir(self, dir, **kwargs):
        self.execute_command('cd {}'.format(dir), **kwargs)

    def touch_file(self, file_name, **kwargs):
        if platform.system() == 'Windows':
            command = 'type NUL >> {file}'
        else:
            command = 'touch {file}'
        logging.debug('Creates file: {}'.format(file_name))
        self.execute_command(command.format(file=file_name), **kwargs)

    def make_dir(self, dir, **kwargs):
        self.execute_command('mkdir {}'.format(dir), **kwargs)

    def _parse_results_file(self, experiment_index, pipeline_collection):
        results = OrderedDict()
        for job_name in experiment_index:
            file_name = pipeline_collection['RESULTS_FILE']
            logging.debug('Reads pipeline results from {}'.format(file_name))
            contents = self.read_file_contents(file_name, directory=str(job_name))
            f_handle = StringIO(contents)
            current_results = pd.Series.from_csv(f_handle)
            results[job_name] = current_results
        return pd.DataFrame(results).T
