from canesm.canesm_database import CanESMensembleDB
from canesm.exceptions import RemoteError
from fabric import Connection
import time


class CanESMsubmitter:

    def __init__(self, db: CanESMensembleDB, delay: int = 0, num_jobs: int = None, debug: bool = False):

        self.db = db
        self.max = num_jobs
        self.n = 0
        self.current_job = None
        self.delay = delay
        self.query_delay = 60  # delay between queries in seconds if waiting on a job setup
        self.debug = debug

    def __iter__(self):

        self.current_job = None
        self.n = 0
        return self

    def __next__(self):

        if self.max is not None and self.n >= self.max:
            raise StopIteration

        try:
            self.current_job = self.db.query('SELECT runid, rundirectory, jobstring '
                                             'FROM status WHERE submitted=0 AND setup=1 LIMIT 1')
        except FileNotFoundError:
            message = 'Could not find the ensemble database. Have you called "setup-ensemble" ' \
                      'before calling "submit-ensemble"?'
            raise FileNotFoundError(message)

        if len(self.current_job) == 0:
            unsubmitted = self.db.query('SELECT runid FROM status WHERE submitted=0 AND setup=0 LIMIT 1')
            if len(unsubmitted) == 0:
                raise StopIteration
            else:
                print('waiting on job setup...')
                self.current_job = False
        else:
            self.n += 1
            self.current_job = self.current_job[0]
        return self

    def submit(self, runid: str = None):
        """
        Submit the next job in the queue to the remote machine

        Parameters
        ----------
        runid:
            job to be submitted on remote machine
        """

        if runid is None:
            runid = self.current_job[0]
            run_directory = self.current_job[1]
            jobstr = self.current_job[2]
        else:
            jobstr = self.db.get(column='jobstring', keys=runid)
            run_directory = self.db.get(column='rundirectory', keys=runid)

        if not runid:
            time.sleep(self.query_delay)
        else:
            with Connection(self.db.machine, self.db.user,
                            gateway=Connection(self.db.gateway_conn, self.db.user)) as remote:
                try:
                    if self.debug:
                        print(f'. env_setup_file && rsub {self.db.machine} {jobstr}')
                    else:
                        with remote.cd(run_directory):
                            remote.run(f'. env_setup_file && rsub {self.db.machine} {jobstr}')
                    self.db.set(column='submitted', keys=runid, values=1)
                    time.sleep(self.delay * 60)
                except RemoteError:
                    raise RemoteError('could not submit job: ' + runid)