import logging
import os
import time
import json

from parsl.channels import LocalChannel
from parsl.launchers import SingleNodeLauncher
from parsl.providers.pbspro.template import template_string
from parsl.providers import TorqueProvider
from parsl.providers.base import JobState, JobStatus

from parsl.providers.torque.torque import translate_table

logger = logging.getLogger(__name__)


class PBSProProvider(TorqueProvider):
    """PBS Pro Execution Provider

    Parameters
    ----------
    channel : Channel
        Channel for accessing this provider. Possible channels include
        :class:`~parsl.channels.LocalChannel` (the default),
        :class:`~parsl.channels.SSHChannel`, or
        :class:`~parsl.channels.SSHInteractiveLoginChannel`.
    account : str
        Account the job will be charged against.
    queue : str
        Queue to request blocks from.
    nodes_per_block : int
        Nodes to provision per block.
    cpus_per_node : int
        CPUs to provision per node.
    init_blocks : int
        Number of blocks to provision at the start of the run. Default is 1.
    min_blocks : int
        Minimum number of blocks to maintain. Default is 0.
    max_blocks : int
        Maximum number of blocks to maintain.
    parallelism : float
        Ratio of provisioned task slots to active tasks. A parallelism value of 1 represents aggressive
        scaling where as many resources as possible are used; parallelism close to 0 represents
        the opposite situation in which as few resources as possible (i.e., min_blocks) are used.
    walltime : str
        Walltime requested per block in HH:MM:SS.
    scheduler_options : str
        String to prepend to the #PBS blocks in the submit script to the scheduler.
    select_options : str
        String to append to the #PBS -l select block in the submit script to the scheduler. This can be used to
        specify ngpus.
    worker_init : str
        Command to be run before starting a worker, such as 'module load Anaconda; source activate env'.
    launcher : Launcher
        Launcher for this provider. The default is
        :class:`~parsl.launchers.SingleNodeLauncher`.
    """
    def __init__(self,
                 channel=LocalChannel(),
                 account=None,
                 queue=None,
                 scheduler_options='',
                 select_options='',
                 worker_init='',
                 nodes_per_block=1,
                 cpus_per_node=1,
                 init_blocks=1,
                 min_blocks=0,
                 max_blocks=1,
                 parallelism=1,
                 launcher=SingleNodeLauncher(),
                 walltime="00:20:00",
                 cmd_timeout=120):
        super().__init__(channel,
                         account,
                         queue,
                         scheduler_options,
                         worker_init,
                         nodes_per_block,
                         init_blocks,
                         min_blocks,
                         max_blocks,
                         parallelism,
                         launcher,
                         walltime,
                         cmd_timeout=cmd_timeout)

        self.template_string = template_string
        self._label = 'pbspro'
        self.cpus_per_node = cpus_per_node
        self.select_options = select_options

    def _status(self):
        '''Returns the status list for a list of job_ids

        Args:
              self

        Returns:
              [status...] : Status list of all jobs
        '''

        job_ids = list(self.resources.keys())
        job_id_list = ' '.join(self.resources.keys())

        jobs_missing = list(self.resources.keys())

        retcode, stdout, stderr = self.execute_wait("qstat -f -F json {0}".format(job_id_list))

        job_statuses = json.loads(stdout)

        if 'Jobs' in job_statuses:
            for job_id, job in job_statuses['Jobs'].items():
                for long_job_id in job_ids:
                    if long_job_id.startswith(job_id):
                        logger.debug('coerced job_id %s -> %s', job_id, long_job_id)
                        job_id = long_job_id
                        break

                job_state = job.get('job_state', JobState.UNKNOWN)
                state = translate_table.get(job_state, JobState.UNKNOWN)
                self.resources[job_id]['status'] = JobStatus(state)
                jobs_missing.remove(job_id)

        # squeue does not report on jobs that are not running. So we are filling in the
        # blanks for missing jobs, we might lose some information about why the jobs failed.
        for missing_job in jobs_missing:
            self.resources[missing_job]['status'] = JobStatus(JobState.COMPLETED)

    def submit(self, command, tasks_per_node, job_name="parsl"):
        """Submits the command job.

        Parameters
        ----------
        command : str
            Command to be executed on the remote side.
        tasks_per_node : int
            Command invocations to be launched per node.
        job_name : str
            Identifier for job.

        Returns
        -------
        None
            If at capacity and cannot provision more
        job_id : str
            Identifier for the job
        """

        job_name = "{0}.{1}".format(job_name, time.time())

        script_path = os.path.abspath("{0}/{1}.submit".format(self.script_dir, job_name))

        logger.debug("Requesting {} nodes_per_block, {} tasks_per_node".format(
            self.nodes_per_block, tasks_per_node)
        )

        job_config = {}
        job_config["submit_script_dir"] = self.channel.script_dir
        job_config["nodes_per_block"] = self.nodes_per_block
        job_config["ncpus"] = self.cpus_per_node
        job_config["walltime"] = self.walltime
        job_config["scheduler_options"] = self.scheduler_options
        job_config["worker_init"] = self.worker_init
        job_config["user_script"] = command

        # Add a colon to select_options if one isn't included
        if self.select_options and not self.select_options.startswith(":"):
            self.select_options = ":" + self.select_options

        job_config["select_options"] = self.select_options

        # Wrap the command
        job_config["user_script"] = self.launcher(command,
                                                  tasks_per_node,
                                                  self.nodes_per_block)

        logger.debug("Writing submit script")
        self._write_submit_script(self.template_string, script_path, job_name, job_config)

        channel_script_path = self.channel.push_file(script_path, self.channel.script_dir)

        submit_options = ''
        if self.queue is not None:
            submit_options = '{0} -q {1}'.format(submit_options, self.queue)
        if self.account is not None:
            submit_options = '{0} -A {1}'.format(submit_options, self.account)

        launch_cmd = "qsub {0} {1}".format(submit_options, channel_script_path)
        retcode, stdout, stderr = self.execute_wait(launch_cmd)

        job_id = None
        if retcode == 0:
            for line in stdout.split('\n'):
                if line.strip():
                    job_id = line.strip()
                    self.resources[job_id] = {'job_id': job_id, 'status': JobStatus(JobState.PENDING)}
        else:
            message = "Command '{}' failed with return code {}".format(launch_cmd, retcode)
            if (stdout is not None) and (stderr is not None):
                message += "\nstderr:{}\nstdout{}".format(stderr.strip(), stdout.strip())
            logger.error(message)

        return job_id

    @property
    def status_polling_interval(self):
        return 60
