
import inspect
import dill as pickle
import tempfile
import os
import sys
import subprocess
import time
from mypyli.samparser import SamRecord

import logging
logging.basicConfig()
LOG = logging.getLogger()
LOG.setLevel("INFO")


class NotFinishedError(ValueError):
    """ Custom error to raise when the job is not finished """


class LSFParallelizer(object):
    """ Parallelizes a command out over multiple nodes """

    def __init__(self):
        pass

    def run_command(self, command):
        pass

    @classmethod
    def wait_for_job(cls, job_name):
        """ waits for job to complete, checks every 10 seconds """
        while cls._job_running(job_name):
            time.sleep(10)

    @staticmethod
    def _job_running(job_name):

        output = subprocess.check_output([
                    "bjobs",
                    "-J", "{job_name}".format(job_name=job_name)
                ])  

        if output:
            return True
        else:
            return False

class Parallelizer(object):
    """ Spawns a function on another LSF node """

    devnull = open("/dev/null", 'w')

    def __init__(self, function, nodes, processors_per_node, imports=[], output_dir="paral_results", job_prefix="job"):

        # set function and optional imports
        self.function = function
        self.imports = imports

        # set some node power values
        self.nodes = nodes
        self.processors_per_node = processors_per_node
    

        self.output_dir = output_dir

        # make directory if doesn't exist -- subject to unlikely race condition, leaving this for simplicity
        if os.path.isdir(output_dir):
            raise ValueError("Output directory '{}' already exists. Please delete it or supply a new output directory.".format(output_dir))
        else:
            os.mkdir(output_dir)

       
        self.jobs = {}
        self.job_prefix = job_prefix
        self.job_id = 1

        
        # pickle the function 
        self.function_pkl = self.output_dir + "/" + "function.pkl"
        with open(self.function_pkl, 'wb') as OUT:
            pickle.dump(function, OUT)


    def run(self, job_args):
        """ 
        Submits a dict of arguments to the cluster to run through the function as soon as possible.
        
        Returns a batch id to use to check status of the job.
        """
        if type(job_args) != dict:
            raise ValueError("job_args argument must be a dict of names parameters to the function supplied on initialization.")

        job_id = self.job_id
        self.job_id += 1
        job_name = self.job_prefix + str(job_id)
            
        args_pkl = self.output_dir + "/" + job_name + "_arguments.pkl"
        results_pkl = self.output_dir + "/" + job_name + "_results.pkl"
        job_script = self.output_dir + "/" + job_name + "_script.py"
        log_base = self.output_dir + "/" + job_name
            
        # pickle the batch
        with open(args_pkl, 'wb') as OUT:
            pickle.dump(job_args, OUT)

        # add the job to the registry
        self.jobs[job_id] = LSFJob(job_id, job_name)
        self.jobs[job_id].results = results_pkl


        # write the script to execute on the compute node
        with open(job_script, 'w') as OUT:
            OUT.write(self.generate_script(args_pkl, results_pkl) + "\n")


        # wait until there is a spot available
        while self.count_jobs_in_queue() >= self.nodes:
            time.sleep(60)

        # run the job
        self.execute_command(job_id, "python {}".format(job_script), log_base)

        return job_id
 
    def get_all_results(self, only_new=False, wait=False):
        """ 
        Yields results of all jobs (or all new jobs)
        
        Optionally blocks until all submitted jobs either complete or fail
        """

        jobs_remaining = self.jobs.values()

        while jobs_remaining:
            jobs_in_next_loop = []
            for job in jobs_remaining:

                # check if job already has finished status and skip if so
                if job.status in ("completed", "exited"):
                    if only_new:
                        continue

                try:
                    yield self.get_results(job.id)
                except NotFinishedError:
                    jobs_in_next_loop.append(job)

            jobs_remaining = jobs_in_next_loop
            jobs_in_next_loop = []

            # block or return
            if wait:
                time.sleep(60)
            else:
                return 




    def get_results(self, job_id, wait=False):
        """ 
        Returns the results of a job or a ValueError if job has failed or a NotFinishedError is job is still in progress
        """

        LOG.debug("Getting result for job id '{}'".format(job_id))

        # make sure the job exists
        try:
            job = self.jobs[job_id]
        except KeyError:
            raise ValueError("No submitted job with id '{}'".format(job_id))


        while True:

            if job.status == "in progress":
                job.update_status()

            # update the status is the previous status was in progress
            if job.status == "in progress":
                if wait:
                    LOG.debug("Waiting for result...")
                    time.sleep(60)
                else:
                    raise NotFinishedError("Job with id '{}' is still in progress.".format(job_id))

            elif job.status == "exited":
                raise ValueError("Job with id '{}' exited with error.".format(job_id))

            elif job.status == "completed":
                with open(job.results, 'rb') as IN:
                    return pickle.load(IN)

            else:
                raise ValueError("Job status '{}' unknown. -- Module error.".format(job.status))


    #
    ## Submitting jobs
    #
    def submit_to_lsf(self, iterable):
        """ Submits jobs up until the maximum node count """
        while self.count_jobs_in_queue() < self.nodes:
            batch = next(iterable)