"""
Note: slurm configuration for debian::

    sudo apt-get install slurm-llnl
    sudo vi /etc/slurm-llnl/slurm.conf
    sudo mkdir /var/run/slurm-llnl
    sudo chmod slurm:slurm /var/run/slurm-llnl

Now put it in rc.d::

    sudo update-rc.d munge defaults
    sudo update-rc.d slurm-llnl defaults

Or run the following each time::

    sudo service munge start
    sudo service slurm-llnl start

"""

import os
import subprocess
import multiprocessing

from .jobid import get_jobid
from . import store, runjob

#from park import config
#from park import environment

# Queue status words
_ACTIVE = ["RUNNING", "COMPLETING"]
_INACTIVE = ["PENDING", "SUSPENDED"]
_ERROR = ["CANCELLED", "FAILED", "TIMEOUT", "NODE_FAIL"]
_COMPLETE = ["COMPLETED"]


class Scheduler(object):
    def jobs(self, status=None):
        """
        Return a list of jobs on the queue.
        """
        #TODO: list completed but not deactivated as well as completed
        #print "queue"
        output,_ = _slurm_exec('squeue','-o', '%i %M %j')
        #print "output",output
        return output

    def deactivate(self, jobid):
        #TODO: remove the job from the active list
        pass

    # Job specific commands
    def submit(self, request, origin):
        """
        Put a command on batch queue, returning its job id.
        """
        #print "submitting job",jobid
        jobid = get_jobid()
        store.create(jobid)
        store.put(jobid,'request',request)

        service = runjob.build_command(jobid,request)

        num_workers = multiprocessing.cpu_count()
        jobdir = store.path(jobid)
        script = os.path.join(jobdir,"J%s.sh"%jobid)
        #commands = ['export %s="%s"'%(k,v) for k,v in config.env().items()]
        commands = ["srun -n 1 -K -o slurm-out.txt nice -n 19 %s &"%service]
#                     "srun -n %d -K -o kernel.out nice -n 19 %s"%(num_workers,kernel)]
        create_batchfile(script,commands)

        _out,err = _slurm_exec('sbatch',
                            '-n',str(num_workers), # Number of tasks
                            #'-K', # Kill if any process returns error
                            #'-o', 'job%j.out',  # output file
                            '-D',jobdir,  # Start directory
                            script)
        if not err.startswith('sbatch: Submitted batch job '):
            raise RuntimeError(err)
        slurmid = err[28:].strip()
        store.put(jobid,'slurmid',slurmid)
        return jobid

    def results(self, id):
        try:
            return runjob.results(id)
        except KeyError:
            pass

        return { 'status': self.status(id) }

    def info(self,id):
        request = store.get(id,'request')
        request['id'] = id
        return request

    def status(self, id):
        """
        Returns the follow states:
        PENDING   --  Job is waiting to be processed
        ACTIVE    --  Job is busy being processed through the queue
        COMPLETE  --  Job has completed successfully
        ERROR     --  Job has either been canceled by the user or an
                      error has been raised
        """

        # Simple checks first
        jobdir = store.path(id)
        if not os.path.exists(jobdir):
            return "UNKNOWN"
        elif store.contains(id, 'results'):
            return "COMPLETE"

        # Translate job id to slurm id
        slurmid = store.get(id,'slurmid')

        # Check slurm queue for job id
        out,_ = _slurm_exec('squeue', '-h', '--format=%i %T')
        out = out.strip()

        state = ''
        inqueue = False
        if out != "":
            for line in out.split('\n'):
                line = line.split()
                if slurmid == line[0]:
                    state = line[1]
                    inqueue = True
                    break

        if inqueue:
            if state in _ACTIVE:
                return "ACTIVE"
            elif state in _INACTIVE:
                return "PENDING"
            elif state in _COMPLETE:
                return "COMPLETE"
            elif state in _ERROR:
                return "ERROR"
            else:
                raise RuntimeError("unexpected state from squeue: %s"%state)
        else:
            return "ERROR"

    def cancel(self, jobid):
        #print "canceling",jobid
        slurmid = store.get(jobid,'slurmid')
        _slurm_exec('scancel',slurmid)

    def delete(self, jobid):
        #jobdir = store.path(jobid)
        self.cancel(jobid)
        store.destroy(jobid)

    def nextjob(self, queue):
        raise NotImplementedError("SLURM queues do not support work sharing")

    def postjob(self, id, results):
        raise NotImplementedError("SLURM queues do not support work sharing")

def create_batchfile(script, commands):
    """
    Create the batchfile to run the job.
    """
    fid = open(script,'w')
    fid.write("#!/bin/sh\n")
    fid.write("\n".join(commands))
    fid.write("\nwait\n")
    fid.close()
    return script

def _slurm_exec(cmd, *args):
    """
    Run a slurm command, capturing any errors.
    """
    #print "cmd",cmd,"args",args
    process = subprocess.Popen([cmd]+list(args),
                               stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)
    out,err = process.communicate()
    if err.startswith(cmd+': error: '):
        raise RuntimeError(cmd+': '+err[15:].strip())
    return out,err