import logging
import os
import re
import time

import typeguard

from parsl.jobs.states import JobState, JobStatus
from parsl.launchers import SingleNodeLauncher
from parsl.launchers.base import Launcher
from parsl.providers.cluster_provider import ClusterProvider
from parsl.providers.condor.template import template_string
from parsl.providers.errors import ScaleOutFailed
from parsl.utils import RepresentationMixin

logger = logging.getLogger(__name__)

from typing import Dict, List, Optional

# See http://pages.cs.wisc.edu/~adesmet/status.html
translate_table = {
    '1': JobState.PENDING,
    '2': JobState.RUNNING,
    '3': JobState.CANCELLED,
    '4': JobState.COMPLETED,
    '5': JobState.FAILED,
    '6': JobState.FAILED,
}


class CondorProvider(RepresentationMixin, ClusterProvider):
    """HTCondor Execution Provider.

    Parameters
    ----------
    nodes_per_block : int
        Nodes to provision per block.
    cores_per_slot : int
        Specify the number of cores to provision per slot. If set to None, executors
        will assume all cores on the node are available for computation. Default is None.
    mem_per_slot : float
        Specify the real memory to provision per slot in GB. If set to None, no
        explicit request to the scheduler will be made. Default is None.
    init_blocks : int
        Number of blocks to provision at time of initialization
    min_blocks : int
        Minimum number of blocks to maintain
    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.
    environment : dict of str
        A dictionary of environmant variable name and value pairs which will be set before
        running a task.
    project : str
        Project which the job will be charged against
    scheduler_options : str
        String to add specific condor attributes to the HTCondor submit script.
    transfer_input_files : list(str)
        List of strings of paths to additional files or directories to transfer to the job
    worker_init : str
        Command to be run before starting a worker.
    requirements : str
        Condor requirements.
    launcher : Launcher
        Launcher for this provider. Possible launchers include
        :class:`~parsl.launchers.SingleNodeLauncher` (the default),
    cmd_timeout : int
        Timeout for commands made to the scheduler in seconds.
    cmd_chunk_size : int
        Calls to the scheduler will be made for chunks of blocks with this size.
    """
    @typeguard.typechecked
    def __init__(self,
                 nodes_per_block: int = 1,
                 cores_per_slot: Optional[int] = None,
                 mem_per_slot: Optional[float] = None,
                 init_blocks: int = 1,
                 min_blocks: int = 0,
                 max_blocks: int = 1,
                 parallelism: float = 1,
                 environment: Optional[Dict[str, str]] = None,
                 project: str = '',
                 scheduler_options: str = '',
                 transfer_input_files: List[str] = [],
                 walltime: str = "00:10:00",
                 worker_init: str = '',
                 launcher: Launcher = SingleNodeLauncher(),
                 requirements: str = '',
                 cmd_timeout: int = 60,
                 cmd_chunk_size: int = 100) -> None:

        label = 'condor'
        super().__init__(label,
                         nodes_per_block,
                         init_blocks,
                         min_blocks,
                         max_blocks,
                         parallelism,
                         walltime,
                         launcher,
                         cmd_timeout=cmd_timeout)
        self.cores_per_slot = cores_per_slot
        self.mem_per_slot = mem_per_slot
        self.cmd_chunk_size = cmd_chunk_size

        # To Parsl, Condor slots should be treated equivalently to nodes
        self.cores_per_node = cores_per_slot
        self.mem_per_node = mem_per_slot

        self.environment = environment if environment is not None else {}
        for key, value in self.environment.items():
            # To escape literal quote marks, double them
            # See: http://research.cs.wisc.edu/htcondor/manual/v8.6/condor_submit.html
            try:
                self.environment[key] = "'{}'".format(value.replace("'", '"').replace('"', '""'))
            except AttributeError:
                pass

        self.project = project
        self.scheduler_options = scheduler_options + '\n'
        self.worker_init = worker_init + '\n'
        self.requirements = requirements
        self.transfer_input_files = transfer_input_files

    def _status(self):
        """Update the resource dictionary with job statuses."""

        for job_id_chunk in _chunker(self.resources.keys(), self.cmd_chunk_size):
            job_id_list = ' '.join(job_id_chunk)
            cmd = "condor_q {0} -af:jr JobStatus".format(job_id_list)
            retcode, stdout, stderr = self.execute_wait(cmd)
            """
            Example output:

            $ condor_q 34524642.0 34524643.0 -af:jr JobStatus
            34524642.0 2
            34524643.0 1
            """

            for line in stdout.splitlines():
                parts = line.strip().split()
                job_id = parts[0]
                state = translate_table.get(parts[1], JobState.UNKNOWN)
                if job_id in self.resources:
                    self.resources[job_id]['status'] = JobStatus(state)

    def status(self, job_ids):
        """Get the status of a list of jobs identified by their ids.

        Parameters
        ----------
        job_ids : list of int
            Identifiers of jobs for which the status will be returned.

        Returns
        -------
        List of int
            Status codes for the requested jobs.

        """
        self._status()
        return [self.resources[jid]['status'] for jid in job_ids]

    def submit(self, command, tasks_per_node, job_name="parsl.condor"):
        """Submits the command onto an Local Resource Manager job.

        example file with the complex case of multiple submits per job:
            Universe =vanilla
            output = out.$(Cluster).$(Process)
            error = err.$(Cluster).$(Process)
            log = log.$(Cluster)
            leave_in_queue = true
            executable = test.sh
            queue 5
            executable = foo
            queue 1

        $ condor_submit test.sub
        Submitting job(s)......
        5 job(s) submitted to cluster 118907.
        1 job(s) submitted to cluster 118908.

        Parameters
        ----------
        command : str
            Command to execute
        job_name : str
            Job name prefix.
        tasks_per_node : int
            command invocations to be launched per node
        Returns
        -------
        None or str
            None if at capacity and cannot provision more; otherwise the identifier for the job.
        """

        logger.debug("Attempting to launch")

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

        scheduler_options = self.scheduler_options
        worker_init = self.worker_init
        if self.mem_per_slot is not None:
            scheduler_options += 'RequestMemory = {}\n'.format(self.mem_per_slot * 1024)
            worker_init += 'export PARSL_MEMORY_GB={}\n'.format(self.mem_per_slot)
        if self.cores_per_slot is not None:
            scheduler_options += 'RequestCpus = {}\n'.format(self.cores_per_slot)
            worker_init += 'export PARSL_CORES={}\n'.format(self.cores_per_slot)

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

        self.environment["JOBNAME"] = "'{}'".format(job_name)

        job_config = {}
        job_config["job_name"] = job_name
        job_config["submit_script_dir"] = self.script_dir
        job_config["project"] = self.project
        job_config["nodes"] = self.nodes_per_block
        job_config["scheduler_options"] = scheduler_options
        job_config["worker_init"] = worker_init
        job_config["user_script"] = command
        job_config["tasks_per_node"] = tasks_per_node
        job_config["requirements"] = self.requirements
        job_config["environment"] = ' '.join(['{}={}'.format(key, value) for key, value in self.environment.items()])

        # Move the user script
        # This is where the command should be wrapped by the launchers.
        wrapped_command = self.launcher(command,
                                        tasks_per_node,
                                        self.nodes_per_block)

        with open(userscript_path, 'w') as f:
            f.write(job_config["worker_init"] + '\n' + wrapped_command)

        the_input_files = [userscript_path] + self.transfer_input_files
        job_config["input_files"] = ','.join(the_input_files)
        job_config["job_script"] = os.path.basename(userscript_path)

        # Construct and move the submit script
        self._write_submit_script(template_string, script_path, job_name, job_config)

        cmd = "condor_submit {0}".format(script_path)
        try:
            retcode, stdout, stderr = self.execute_wait(cmd)
        except Exception as e:
            raise ScaleOutFailed(self.label, str(e))

        job_id = []

        if retcode == 0:
            for line in stdout.split('\n'):
                if re.match('^[0-9]', line) is not None:
                    cluster = line.split(" ")[5]
                    # We know the first job id ("process" in condor terms) within a
                    # cluster is 0 and we know the total number of jobs from
                    # condor_submit, so we use some list comprehensions to expand
                    # the condor_submit output into job IDs
                    # e.g., ['118907.0', '118907.1', '118907.2', '118907.3', '118907.4', '118908.0']
                    processes = [str(x) for x in range(0, int(line[0]))]
                    job_id += [cluster + process for process in processes]

            self._add_resource(job_id)
            return job_id[0]
        else:
            message = "Command '{}' failed with return code {}".format(cmd, retcode)
            message += " and standard output '{}'".format(stdout.strip()) if stdout is not None else ''
            message += " and standard error '{}'".format(stderr.strip()) if stderr is not None else ''
            raise ScaleOutFailed(self.label, message)

    def cancel(self, job_ids):
        """Cancels the jobs specified by a list of job IDs.

        Parameters
        ----------
        job_ids : list of str
            The job IDs to cancel.

        Returns
        -------
        list of bool
            Each entry in the list will be True if the job is cancelled succesfully, otherwise False.
        """
        rets = []
        for job_id_chunk in _chunker(job_ids, self.cmd_chunk_size):
            job_id_list = ' '.join(job_id_chunk)
            cmd = "condor_rm {0}; condor_rm -forcex {0}".format(job_id_list)
            logger.debug("Attempting removal of jobs : {0}".format(cmd))
            retcode, stdout, stderr = self.execute_wait(cmd)
            if retcode == 0:
                for jid in job_id_chunk:
                    if jid in self.resources:
                        self.resources[jid]['status'] = JobStatus(JobState.CANCELLED)
                rets.extend([True for i in job_id_chunk])
            else:
                rets.extend([False for i in job_id_chunk])

        return rets

    def _add_resource(self, job_id):
        for jid in job_id:
            self.resources[jid] = {'status': JobStatus(JobState.PENDING), 'size': 1}
        return True

    @property
    def status_polling_interval(self):
        return 60


# see https://stackoverflow.com/a/18243990
def _chunker(seq, size):
    res = []
    for el in seq:
        res.append(el)
        if len(res) == size:
            yield res
            res = []
    if res:
        yield res


if __name__ == "__main__":

    print("None")