from canesm.util import RemoteDBConn
from fabric import Connection
from typing import List, Union
import pandas as pd
from canesm.canesm_setup import CanESMsetup


class CanESMensembleDB:
    """
    Database to hold the ensemble status.

    Parameters
    ----------
    db_file:
        name of database file
    machine:
        name of remote machine
    user:
        account on remote machine used for ssh connection
    gateway_conn:
        ssh gateway

    """
    def __init__(self, db_file: str, machine: str, user: str, gateway_conn: str):

        self.db_file = db_file

        self.machine = machine
        self.user = user
        self.gateway_conn = gateway_conn
        self.conn = RemoteDBConn(self.db_file, self.machine, self.user, self.gateway_conn)

    def get(self, column: str, keys: Union[str, List[str]]) -> Union[List[Union[str, int]], str, int]:
        """
        Get the values in `column` associated with the primary keys.

        Parameters
        ----------
        column:
            column name to be queried
        keys:
            primary keys

        Returns
        -------
            values from `column`. If a list of keys is provided this will be a List, if a value is provided a single
            value is returned
        """
        values = []
        with self.conn as conn:
            if type(keys) is not list:
                keys = [keys]
            for key in keys:
                c = conn.cursor()
                c.execute(f'SELECT {column} FROM status WHERE runid=? ORDER BY priority', (key,))
                values.append(c.fetchall()[0][0])

        if len(values) == 1:
            return values[0]
        else:
            return values

    def set(self, column: str, keys: Union[str, List[str]], values: Union[List[Union[str, int]], str, int]):
        """
        Set the values in `column` associated with the primary keys.

        Parameters
        ----------
        column:
            column name to be queried
        keys:
            primary keys
        values:
            values to be set
        """
        with self.conn as conn:
            if type(values) is not list:
                values = [values]
            if type(keys) is not list:
                keys = [keys]
            for key, value in zip(keys, values):
                c = conn.cursor()
                c.execute(f'UPDATE status SET {column}=? WHERE runid=?', (value, key))
                conn.commit()

        return self

    def query(self, query: str):
        """
        Execute a query on the ensemble database

        Parameters
        ----------
        query:
            string to be executed

        Returns
        -------
            If `query` contains `select` the values are returned, otherwise None
        """
        with self.conn as conn:
            c = conn.cursor()
            c.execute(query)

            if 'select' in query.lower():
                return c.fetchall()

            if 'update' in query.lower():
                conn.commit()

        return self

    def show_table(self):
        """
        Print the current state of the ensemble database
        """
        with self.conn as conn:
            print(pd.read_sql_query('SELECT runid, submitted, setup FROM status '
                                    'ORDER BY submitted DESC', conn).to_string(index=False))

    def setup(self, ensemble_jobs: List[CanESMsetup]):
        """
        Setup a database to store whether the ensemble members have been submitted. If a database for the ensemble
        already exists then entries will be added to it.
        """
        with self.conn as conn:

            c = conn.cursor()
            c.execute('''CREATE TABLE IF NOT EXISTS 
                         status (runid TEXT PRIMARY KEY,
                         submitted INTEGER DEFAULT 9, 
                         jobidx INTEGER,
                         jobstring TEXT DEFAULT "",
                         rundirectory TEXT,
                         priority INTEGER DEFAULT 1,
                         setup INTEGER DEFAULT 0)''')

            jobs = []
            for idx, job in enumerate(ensemble_jobs):
                jobs.append((job.runid, 0, idx, job.job_str, job.run_directory, idx, 0))
            c.executemany('''INSERT OR IGNORE INTO status VALUES (?,?,?,?,?,?,?)''', jobs)
            conn.commit()

        return self

    def reset(self):
        """
        reset the submitted and setup status of the database
        """

        with self.conn as conn:
            c = conn.cursor()
            c.execute('UPDATE status SET submitted=0')
            c.execute('UPDATE status SET setup=0')
            conn.commit()

        return self

    def delete_dbfile(self):
        """
        delete the database
        """

        with Connection(self.machine, self.user, gateway=Connection(self.gateway_conn, self.user)) as c:
            c.run('rm ' + self.db_file)

        return self