'''
/***************************************************************************
Name	   :  datasets_model.py
Description:  Datasets model class for FEWSTools plugin,
              updated from QGIS2
copyright  :  (C) 2019-2023 by FEWS
email      :  minxuansun@contractor.usgs.gov
Modified   :  11/19/2019 - cholen - Updated
              unset_default_flag_on_non_selected_datasetsDescription
              02/13/2020 - cholen - Added method to check names
              12/28/2020 - cholen - Add datascale factor
              07/22/2022 - cholen - Add migrations
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/
'''
import copy
import os

from PyQt5.QtWidgets import QMessageBox

from fews_tools import fews_tools_config as config
from fews_tools.utilities import database_utilities as db_util


class DatasetsModel:
    '''
    Class to hold dataset table record information, methods
    datasets_attributes = [
            'name', 'type', 'extent', 'PERIODICITY', 'folder',
            'prefix', 'dateformat', 'suffix', 'nodata',
            'avgfolder', 'avgprefix', 'avgdateformat','avgsuffix',
            'avgnodata','REMOTEHOST','remotedir','usr','pw','zipdate',
            'zipprefix','zipsuffix','projection',
            'ucdateformat','ucprefix','uxsuffix','default']
    '''
    def __init__(self):
        self.datasets_table_name = 'Datasets'
        self.ds_dic = None

        self.update_sql = (
            "UPDATE " + self.datasets_table_name +
            " SET DATATYPE=?,DATAEXTENT=?, "
            "PERIODICITY=?,DATAFOLDER=?,DATAPREFIX=?,DATADATEFORMAT=?,"
            "DATASUFFIX=?,DATAMISSINGVALUE=?,AVGDATAFOLDER=?,"
            "AVGDATAPREFIX=?,AVGDATADATEFORMAT=?,AVGDATASUFFIX=?,"
            "AVGDATAMISSINGVALUE=?,REMOTEHOST=?,REMOTEDIRECTORY=?,"
            "USERNAME=?,PASSWORD=?,FILEDATEFORMAT=?,FILEPREFIX=?,"
            "FILESUFFIX=?,UNCOMPRESSEDDATAPROJECTION=?,"
            "UNCOMPRESSEDDATADATEFORMAT=?,UNCOMPRESSEDDATAPREFIX=?,"
            "UNCOMPRESSEDDATASUFFIX=?,DEFAULTDS=?,"
            "DATASCALEFACTOR=? WHERE DATASETNAME = ?")
        self.insert_sql = (
            "INSERT INTO " + self.datasets_table_name +
            " (DATASETNAME,DATATYPE,DATAEXTENT,PERIODICITY,"
            " DATAFOLDER, DATAPREFIX, DATADATEFORMAT,DATASUFFIX,"
            " DATAMISSINGVALUE,AVGDATAFOLDER,AVGDATAPREFIX,"
            " AVGDATADATEFORMAT,AVGDATASUFFIX,AVGDATAMISSINGVALUE,"
            " REMOTEHOST,REMOTEDIRECTORY,USERNAME,PASSWORD,"
            " FILEDATEFORMAT,FILEPREFIX,FILESUFFIX,"
            " UNCOMPRESSEDDATAPROJECTION,UNCOMPRESSEDDATADATEFORMAT,"
            " UNCOMPRESSEDDATAPREFIX,UNCOMPRESSEDDATASUFFIX,"
            " DEFAULTDS,DATASCALEFACTOR)"
            " VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,"
            " ?,?,?,?,?,?,?,?,?,?,?)")

        initial_default_dataset = self.query_default_dataset_name()
        self.query_named_dataset(initial_default_dataset)


    def delete_dataset_record(self, ds_name):
        '''
        Deletes selected dataset

        :param(string) ds_name: Name of dataset to delete
        '''
        sql = ('DELETE FROM ' + self.datasets_table_name +
               ' WHERE DATASETNAME = ?')
        params = (ds_name, )
        db_util.write_db(sql, params)

    def get_all_dataset_names(self, db_name=None):
        '''
        Extract dataset names from datasets table. Used to fill
        dataset combo box when loading form.

        :param(string)db_name: Database name (default=None, will use config)
        :return(list)dataset_list: Dataset name list
        '''
        if db_name is None:
            db_name = config.DB_FILE
        sql = 'SELECT * FROM ' + self.datasets_table_name
        qry_result = db_util.query_db(sql, None, True, db_name)
        dataset_list = []
        for row in qry_result:
            dataset_list.append(row[0])
        return dataset_list

    def get_ds_dictionary(self):
        '''
        Return dataset dictionary
        '''
        return self.ds_dic

    def query_named_dataset(self, ds_name, db_name=None):
        '''
        Extracts dataset information from database and loads it into the class
        variables.
        :param(string) ds_name: Dataset name.
        :param(string) db_name: Database name (default=None, will use config)
        :return err(boolean): True if not found, else false
        '''
        if db_name is None:
            db_name = config.DB_FILE
        err = True
        sql = ("SELECT * FROM " + self.datasets_table_name +
               " WHERE DATASETNAME=?")
        params = (str(ds_name), )
        qry_result = db_util.query_db(sql, params, False, db_name)
        if qry_result:  # update the dictionary with the returned query
            self.ds_dic = {
                'DATASETNAME': qry_result[0],
                'DATATYPE': qry_result[1],
                'DATAEXTENT': qry_result[2],
                'PERIODICITY': qry_result[3],
                'DATAFOLDER': qry_result[4],
                'DATAPREFIX': qry_result[5],
                'DATADATEFORMAT': qry_result[6],
                'DATASUFFIX': qry_result[7],
                'DATAMISSINGVALUE': qry_result[8],
                'AVGDATAFOLDER': qry_result[9],
                'AVGDATAPREFIX': qry_result[10],
                'AVGDATADATEFORMAT': qry_result[11],
                'AVGDATASUFFIX': qry_result[12],
                'AVGDATAMISSINGVALUE': qry_result[13],
                'REMOTEHOST': qry_result[14],
                'REMOTEDIRECTORY': qry_result[15],
                'USERNAME': qry_result[16],
                'PASSWORD': qry_result[17],
                'FILEDATEFORMAT': qry_result[18],
                'FILEPREFIX': qry_result[19],
                'FILESUFFIX': qry_result[20],
                'UNCOMPRESSEDDATAPROJECTION': qry_result[21],
                'UNCOMPRESSEDDATADATEFORMAT': qry_result[22],
                'UNCOMPRESSEDDATAPREFIX': qry_result[23],
                'UNCOMPRESSEDDATASUFFIX': qry_result[24],
                'DEFAULTDS': qry_result[25]}
            # in old db this didn't exist, need to handle the migration
            try:
                self.ds_dic['DATASCALEFACTOR'] = qry_result[26]
            except IndexError:
                self.ds_dic['DATASCALEFACTOR'] = 1

            err = False
        return err

    def query_default_dataset_name(self, ds_type=config.DATA_TYPES[0][1]):
        '''
        Extracts default dataset name and returns it
        If none are marked as default it returns the first dataset record
        :param ds_type: Dataset type. Default is PPT (Rainfall data)
        :return ret_val(string): Default dataset name
        '''
        sql = ("SELECT * FROM " + self.datasets_table_name +
               " WHERE DEFAULTDS = 'Yes' AND DATATYPE = '" + ds_type + "'")
        qry_result = db_util.query_db(sql)
        if qry_result:
            ret_val = qry_result[0]
        else:
            # get the first dataset in the list
            ds_name_list = self.get_all_dataset_names()
            ret_val = ds_name_list[0]
        return ret_val

    def update_dataset_filepaths(self, workspace):
        '''
        Function to update the datasets table with the workspace value.
        :params workspace(string) - Workspace path
        '''
        # fix the special case of having selected the root
        if workspace == os.path.abspath(os.sep):
            workspace = 'C:'
        # handle any case where we don't end with the os.sep
        if workspace[-1] != os.sep:
            workspace += os.sep

        sql = 'SELECT * FROM ' + self.datasets_table_name
        qry_result = db_util.query_db(sql, None, True)
        trim_text = str(qry_result[0][4]).split('ProgramSettings')[0]
        params = (trim_text, workspace, trim_text, workspace)
        sql_write = ('UPDATE ' + self.datasets_table_name + ' SET '
                     'DATAFOLDER = REPLACE (DATAFOLDER, ?, ?), '
                     'AVGDATAFOLDER = REPLACE (AVGDATAFOLDER, ?, ?)')
        db_util.write_db(sql_write, params)

    def unset_default_flag_on_non_selected_datasets(self, default_ds_name):
        '''
        Update the default dataset.
        :params default_ds_name(string): The new default dataset name.
        '''
        # Get default dataset's data type
        self.query_named_dataset(default_ds_name)
        data_type = self.get_ds_dictionary()['DATATYPE']
        # Update dataset
        params = (default_ds_name, data_type)
        sql_write = \
            ("UPDATE " + self.datasets_table_name +
             " SET DEFAULTDS = 'No' WHERE DATASETNAME != ? AND DATATYPE = ?")
        db_util.write_db(sql_write, params)

        sql_write = \
            ("UPDATE " + self.datasets_table_name +
             " SET DEFAULTDS = 'Yes' WHERE DATASETNAME = ? AND DATATYPE = ?")
        db_util.write_db(sql_write, params)

    def write_dataset_record(self, widget_vals):
        '''
        Writes or updates a row in the datasets table.
        :params
            widget_vals(list): The widget values from form in this order.
                'DATASETNAME,DATATYPE,DATAEXTENT,PERIODICITY,
                'DATAFOLDER, DATAPREFIX, DATADATEFORMAT,DATASUFFIX,
                'DATAMISSINGVALUE,AVGDATAFOLDER,AVGDATAPREFIX,
                'AVGDATADATEFORMAT,AVGDATASUFFIX,AVGDATAMISSINGVALUE,
                'REMOTEHOST,REMOTEDIRECTORY,USERNAME,PASSWORD,
                'FILEDATEFORMAT,FILEPREFIX,FILESUFFIX,
                'UNCOMPRESSEDDATAPROJECTION,UNCOMPRESSEDDATADATEFORMAT,
                'UNCOMPRESSEDDATAPREFIX,UNCOMPRESSEDDATASUFFIX,
                'DEFAULTDS','DATASCALEFACTOR'
        '''
        err = self.query_named_dataset(widget_vals[0])
        loc_widget_vals = copy.deepcopy(widget_vals)
        if err is False:
            loc_widget_vals.append(loc_widget_vals.pop(0))
            sql_write = self.update_sql
        else:
            sql_write =self.insert_sql

        params = tuple(loc_widget_vals)
        db_util.write_db(sql_write, params)
        # update the default dataset if necessary
        if widget_vals[-2] == 'Yes':
            self.unset_default_flag_on_non_selected_datasets(widget_vals[0])

    def write_dataset_record_from_dic(self, ds_dic):
        '''
        Writes or updates a row in the datasets table using a dataset dictionary.
        :params
            reg_dic: The 'widget' values need to be order.
        '''
        param_list = [
                ds_dic['DATASETNAME'],
                ds_dic['DATATYPE'],
                ds_dic['DATAEXTENT'],
                ds_dic['PERIODICITY'],
                ds_dic['DATAFOLDER'],
                ds_dic['DATAPREFIX'],
                ds_dic['DATADATEFORMAT'],
                ds_dic['DATASUFFIX'],
                ds_dic['DATAMISSINGVALUE'],
                ds_dic['AVGDATAFOLDER'],
                ds_dic['AVGDATAPREFIX'],
                ds_dic['AVGDATADATEFORMAT'],
                ds_dic['AVGDATASUFFIX'],
                ds_dic['AVGDATAMISSINGVALUE'],
                ds_dic['REMOTEHOST'],
                ds_dic['REMOTEDIRECTORY'],
                ds_dic['USERNAME'],
                ds_dic['PASSWORD'],
                ds_dic['FILEDATEFORMAT'],
                ds_dic['FILEPREFIX'],
                ds_dic['FILESUFFIX'],
                ds_dic['UNCOMPRESSEDDATAPROJECTION'],
                ds_dic['UNCOMPRESSEDDATADATEFORMAT'],
                ds_dic['UNCOMPRESSEDDATAPREFIX'],
                ds_dic['UNCOMPRESSEDDATASUFFIX'],
                ds_dic['DEFAULTDS'],
                ds_dic['DATASCALEFACTOR']
                ]
        self.write_dataset_record(param_list)

    def check_new_dataset_name(self, new_name):
        '''
        Check if name exists return err if yes
        :param new_name(string) - New name
        :return err(boolean) - True if name already exists, else False
        '''
        err = False
        dataset_name_list = self.get_all_dataset_names()
        if new_name in dataset_name_list:
            err = True
            QMessageBox.critical(
                None, 'Error - Name already exists',
                'The selected name already exists in the database.',
                QMessageBox.Ok)
        return err

    def filter_by_data_type(self, data_type: str) -> dict:
        '''
        Get list of dictionary that contains the datasets info whose data
        type match the given data type
        :param data_type(string): A string specify the data type (3 letters)
        :return return_value(dict): A dict. Each key is the dataset name,
        the relative value is a dict contains the dataset_info
        '''
        return_value = {}
        dataset_name_list = self.get_all_dataset_names()
        for name in dataset_name_list:
            self.query_named_dataset(name)
            if self.ds_dic['DATATYPE'] == data_type.upper():
                # only stores the dataset info if its data type matches
                # the given one
                return_value[name] = self.ds_dic
        return return_value

    def add_data_scale_factor_field_to_table(self):
        '''
        Needs to run on startup to add column if it isn't there
        later can remove this function when all users have this column in their
        database
        '''
        col_names = self.get_current_columns_from_table()
        if 'DATASCALEFACTOR' not in col_names:
            sql_write = ("ALTER TABLE " + self.datasets_table_name +
                         " ADD COLUMN 'DATASCALEFACTOR' INT")
            db_util.write_db(sql_write)
        # recheck
        col_names = self.get_current_columns_from_table()
        if 'DATASCALEFACTOR' in col_names:
            sql_write = \
                ("UPDATE " + self.datasets_table_name +
                 " SET DATASCALEFACTOR = 1 WHERE DATATYPE != 'PET'")
            db_util.write_db(sql_write)

            sql_write = \
                ("UPDATE " + self.datasets_table_name +
                 " SET DATASCALEFACTOR = 100 WHERE DATATYPE = 'PET'")
            db_util.write_db(sql_write)

    def get_current_columns_from_table(self):
        '''
        Get the current columns from the table
        '''
        sql_query = ("SELECT * FROM " + self.datasets_table_name)
        col_names = db_util.query_db_columns(sql_query)
        return col_names

    def migrate_datasets(self):
        '''
        Migrate the datasets rows to the new database
        '''
        # get old db datasets
        old_ds_list = self.get_all_dataset_names(config.OLD_DB_FILE)
        new_ds_list = self.get_all_dataset_names(config.DB_FILE)
        for entry in old_ds_list:
            if entry in new_ds_list:  # should be up to date
                continue
            # get old record
            err = self.query_named_dataset(entry, config.OLD_DB_FILE)
            ds_dic = self.get_ds_dictionary()
            self.write_dataset_record_from_dic(ds_dic)
