import pytest
from fabric import Connection
from invoke.exceptions import UnexpectedExit
import sys
import os

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

config_file = 'db_test.yaml'
setup_folder = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'setup_files')
from canesm import CanESMensemble
from canesm import CanESMsubmitter
from canesm import CanESMensembleDB
from canesm.util import RemoteDBConn


class DB(object):
    def __init__(self):
        self.ens = CanESMensemble.from_config_file(os.path.join(setup_folder, config_file))

    def setup_db(self):
        with Connection(self.ens.machine, user=self.ens.user,
                        gateway=Connection(self.ens.gateway_conn, user=self.ens.user)) as c:
            c.run('mkdir -p ' + self.ens.run_directory[0])
            try:
                c.run('rm ' + self.ens.db.db_file)
            except UnexpectedExit:
                pass

        canjob = CanESMensembleDB(self.ens.db.db_file, self.ens.machine, self.ens.user, self.ens.gateway_conn)
        canjob.setup(self.ens.jobs)

        runids = [job.runid for job in self.ens.jobs]
        # fake the setup
        for runid in runids:
            canjob.set('setup', runid, 1)

    def teardown_db(self):
        with Connection(self.ens.machine, user=self.ens.user,
                        gateway=Connection(self.ens.gateway_conn, user=self.ens.user)) as c:
            c.run('rm -rf ' + self.ens.run_directory[0])


@pytest.fixture(scope="module")
def db():
    return DB()


class TestSubmitJob:

    @pytest.fixture(autouse=True)
    def db_setup(self, db):
        db.setup_db()
        yield
        db.teardown_db()

    def test_submit_job(self):

        ensemble = CanESMensemble.from_config_file(os.path.join(setup_folder, config_file))

        canjob = CanESMsubmitter(ensemble.db, debug=True)
        canjob.submit('run-001a')
        canjob.submit('run-003a')

        with RemoteDBConn(ensemble.db.db_file, ensemble.machine, ensemble.user, ensemble.gateway_conn) as conn:
            c = conn.cursor()
            c.execute('''SELECT * FROM status''')
            test = c.fetchall()

        assert test == [('run-001a', 1, 0, 'run-001a_1990m01_2014m12_hare_job',
                         'canesm_ensemble/testing/db_test/run-001a', 0, 1),
                        ('run-002a', 0, 1, 'run-002a_1990m01_2014m12_hare_job',
                         'canesm_ensemble/testing/db_test/run-002a', 1, 1),
                        ('run-003a', 1, 2, 'run-003a_1990m01_2014m12_hare_job',
                         'canesm_ensemble/testing/db_test/run-003a', 2, 1),
                        ('run-004a', 0, 3, 'run-004a_1990m01_2014m12_hare_job',
                         'canesm_ensemble/testing/db_test/run-004a', 3, 1),
                        ('run-005a', 0, 4, 'run-005a_1990m01_2014m12_hare_job',
                         'canesm_ensemble/testing/db_test/run-005a', 4, 1)]

    def test_full_cycle(self):

        ensemble = CanESMensemble.from_config_file(os.path.join(setup_folder, config_file))

        for job in CanESMsubmitter(ensemble.db, delay=0, debug=True):
            job.submit()

        with RemoteDBConn(ensemble.db.db_file, ensemble.machine, ensemble.user, ensemble.gateway_conn) as conn:
            c = conn.cursor()
            c.execute('''SELECT * FROM status''')
            test = c.fetchall()

        assert test == [('run-001a', 1, 0, 'run-001a_1990m01_2014m12_hare_job',
                         'canesm_ensemble/testing/db_test/run-001a', 0, 1),
                        ('run-002a', 1, 1, 'run-002a_1990m01_2014m12_hare_job',
                         'canesm_ensemble/testing/db_test/run-002a', 1, 1),
                        ('run-003a', 1, 2, 'run-003a_1990m01_2014m12_hare_job',
                         'canesm_ensemble/testing/db_test/run-003a', 2, 1),
                        ('run-004a', 1, 3, 'run-004a_1990m01_2014m12_hare_job',
                         'canesm_ensemble/testing/db_test/run-004a', 3, 1),
                        ('run-005a', 1, 4, 'run-005a_1990m01_2014m12_hare_job',
                         'canesm_ensemble/testing/db_test/run-005a', 4, 1)]

    def test_njobs(self):
        num_jobs = 3
        ensemble = CanESMensemble.from_config_file(os.path.join(setup_folder, config_file))

        for job in CanESMsubmitter(ensemble.db, delay=0, num_jobs=num_jobs, debug=True):
            job.submit()

        with RemoteDBConn(ensemble.db.db_file, ensemble.machine, ensemble.user, ensemble.gateway_conn) as conn:
            c = conn.cursor()
            c.execute('''SELECT * FROM status''')
            test = c.fetchall()

        assert test == [('run-001a', 1, 0, 'run-001a_1990m01_2014m12_hare_job',
                         'canesm_ensemble/testing/db_test/run-001a', 0, 1),
                        ('run-002a', 1, 1, 'run-002a_1990m01_2014m12_hare_job',
                         'canesm_ensemble/testing/db_test/run-002a', 1, 1),
                        ('run-003a', 1, 2, 'run-003a_1990m01_2014m12_hare_job',
                         'canesm_ensemble/testing/db_test/run-003a', 2, 1),
                        ('run-004a', 0, 3, 'run-004a_1990m01_2014m12_hare_job',
                         'canesm_ensemble/testing/db_test/run-004a', 3, 1),
                        ('run-005a', 0, 4, 'run-005a_1990m01_2014m12_hare_job',
                         'canesm_ensemble/testing/db_test/run-005a', 4, 1)]