from typing import Self
from psycopg2 import sql as pgsql
from qgis.PyQt.QtCore import Qt, QCoreApplication
from qgis.core import QgsTask
from qgis.PyQt.QtWidgets import QInputDialog, QMessageBox, QListWidgetItem, QPushButton
# Import the code for the dialog
try:
    from widgets.table_managment_dialog import TableMgmtDialog
    from support_scripts.__init__ import TR
    from support_scripts.qt_data import _check_state, _item_flag, _match_flag
except ModuleNotFoundError:
    from ..widgets.table_managment_dialog import TableMgmtDialog
    from ..support_scripts.__init__ import TR
    from ..support_scripts.qt_data import _check_state, _item_flag, _match_flag
from ..support_scripts.notifier import report_warning
from ..support_scripts.notifier import log as gdf_log
__author__ = 'Axel Horteborn'


class TableManagement:
    def __init__(self: Self, parent) -> None:
        """A widget that enables the possibility to insert data from a text
        file into a shapefile"""
        self.add_to_Param_row_count = 0
        self.db = parent.db
        self.dock_widget = parent.dock_widget
        self.parent = parent
        translate = TR('TableManagement')
        self.tr = translate.tr
        # Create the dialog (after translation) and keep reference
        self.TMD = TableMgmtDialog()
        self.tables_in_db = 0
        self.params_in_list = 0
        self.current_table = None
        self.current_schema = None
        self.params_in_table = None
        self.items_in_table = []

    def run(self: Self) -> None:
        """Connects the push buttons and enable the visibility of the dialog."""
        self.TMD.pButRemove.clicked.connect(self.remove_table_from_db)
        self.TMD.pButCombine.clicked.connect(self.merge_tbls)
        self.TMD.pButChangeTbl.clicked.connect(self.edit_tbl_name)
        self.TMD.pButChangeParam.clicked.connect(self.edit_param_name)
        self.TMD.pButAdd_Param.clicked.connect(self.retrieve_params)
        self.TMD.pButSave.clicked.connect(self.save_table)
        self.TMD.pButGetYieldCol.clicked.connect(self.update_column_list)
        self.TMD.pButSplitRows.clicked.connect(self.split_rows)
        self.TMD.PBMakeRows.clicked.connect(self.make_rows)
        self.update_table_list()
        self.TMD.show()
        if not self.parent.test_mode:
            self.TMD.exec()

    def merge_tbls(self: Self) -> None:
        """Merging two data sets into one."""
        tables_to_merge = []
        new_name = self.TMD.LEName.text()
        new_type = self.TMD.CBDataType.currentText()
        new_schema = None
        if new_type == self.tr('plant'):
            new_schema = 'plant'
        if new_type == self.tr('fertilize'):
            new_schema = 'ferti'
        if new_type == self.tr('spray'):
            new_schema = 'spray'
        if new_type == self.tr('other'):
            new_schema = 'other'
        if new_type == self.tr('harvest'):
            new_schema = 'harvest'
        if new_type == self.tr('soil'):
            new_schema = 'soil'
        if new_name == '':
            report_warning(self.tr('You need to fill in a new name'))
            return
        if new_schema is None:
            report_warning(self.tr('You have to decide what type of data it is'))
            return
        if new_name in self.db.get_tables_in_db(new_schema):
            report_warning(self.tr('You need a new name'))
            return
        c = 0
        for item in self.items_in_table:
            if item.checkState() == 2:
                c += 1
                tables_to_merge.append(item.text())
        if c < 2:
            report_warning(self.tr('You need at least 2 dataset when merging'))
            return
        union_parts = []
        for tbl_full in tables_to_merge:
            parts = tbl_full.split('.')
            if len(parts) == 2:
                tbl_id = pgsql.SQL("{}.{}").format(
                    pgsql.Identifier(parts[0]), pgsql.Identifier(parts[1]))
            else:
                tbl_id = pgsql.Identifier(tbl_full)
            union_parts.append(pgsql.SQL("SELECT * FROM {tbl}").format(tbl=tbl_id))
        query = pgsql.SQL("CREATE TABLE {schema}.{new} AS ({body})").format(
            schema=pgsql.Identifier(new_schema),
            new=pgsql.Identifier(new_name),
            body=pgsql.SQL(" UNION ").join(union_parts))
        self.db.execute_sql(query)
        self.db.update_row_id(new_schema, new_name)
        self.db.create_indexes(new_name, [], new_schema)
        self.TMD.LEName.setText('')
        self.TMD.CBDataType.setCurrentIndex(0)
        self.update_table_list()

    def check_multiple(self: Self) -> tuple[bool, str]:
        """Checks if multiple table is selected

        Returns
        -------
        Bool, False if more than one row is selected
        str, schema.table_name"""
        c = 0
        for item in self.items_in_table:
            if item.checkState() == 2:
                c += 1
                table = item.text()
        if c != 1:
            report_warning(self.tr('You can only have one dataset selected'))
            return False, 'failed'
        return True, table

    def retrieve_params(self):
        """This function is trigged by "edit" table, and basically fills the left list widget,
        all attributes that have a index get selected by default."""
        suc, table = self.check_multiple()
        if suc is False:
            return
        self.current_schema, self.current_table = table.split('.')
        if self.params_in_list != 0:
            model = self.TMD.SAParams.model()
            for item in self.params_in_table:
                qIndex = self.TMD.SAParams.indexFromItem(item)
                model.removeRow(qIndex.row())
        self.params_in_list = 0
        checked_params = []
        table = self.current_table
        schema = self.current_schema
        indexes = self.db.get_indexes(table, schema)
        for nbr in indexes.keys():
            checked_params.append(indexes[nbr]['index_col'])
        columns = self.db.get_all_columns(table, schema)
        for param_name in columns:
            if param_name in ['cmin', 'xmin', 'xmax', 'cmax', 'ctid', 'pos',
                                 'polygon', 'tableoid', '_', 'field_row_id']:
                continue
            if param_name[:3] == '...':
                continue
            item_name = str(param_name)
            testcase_name = QCoreApplication.translate("qadashboard", item_name, None)
            item = QListWidgetItem(testcase_name, self.TMD.SAParams)
            item.setFlags(item.flags() | _item_flag('ItemIsUserCheckable'))
            if param_name in checked_params:
                item.setCheckState(_check_state('Checked'))
            else:
                item.setCheckState(_check_state('Unchecked'))
            self.params_in_list += 1
        self.params_in_table = self.TMD.SAParams.findItems('', _match_flag('MatchContains'))

    def save_table(self):
        """Updates the attribute indexes that are checked in the list widget"""
        checked_params = []
        table = self.current_table
        schema = self.current_schema
        indexes = self.db.get_indexes(table, schema)
        create_index_for = []
        remove_index_for = []
        for nbr in indexes.keys():
            checked_params.append(indexes[nbr]['index_col'])
        for item in (self.params_in_table or []):
            if item.text() in checked_params:
                if item.checkState() == 0:
                    remove_index_for.append(item.text())
            else:
                if item.checkState() == 2:
                    create_index_for.append(item.text())
        for index in create_index_for:
            self.db.execute_sql(
                pgsql.SQL("CREATE INDEX {idx} ON {schema}.{tbl} USING btree({col})").format(
                    idx=pgsql.Identifier(f"{index}_{schema}_{table}"),
                    schema=pgsql.Identifier(schema),
                    tbl=pgsql.Identifier(table),
                    col=pgsql.Identifier(index)))
        for index in remove_index_for:
            self.db.execute_sql(
                pgsql.SQL("DROP INDEX IF EXISTS {schema}.{idx}").format(
                    schema=pgsql.Identifier(schema),
                    idx=pgsql.Identifier(f"{index}_{schema}_{table}")))
        if schema != 'harvest':
            self.db.execute_sql(
                pgsql.SQL("DROP INDEX IF EXISTS {schema}.{idx};"
                          " CREATE INDEX {idx} ON {schema}.{tbl} USING gist(polygon)").format(
                    schema=pgsql.Identifier(schema),
                    idx=pgsql.Identifier(f"gist_{table}"),
                    tbl=pgsql.Identifier(table)))
        try:
            if schema != 'weather':
                self.db.execute_sql(
                    pgsql.SQL("DROP INDEX IF EXISTS {schema}.{idx};"
                              " CREATE INDEX {idx} ON {schema}.{tbl} USING gist(pos)").format(
                        schema=pgsql.Identifier(schema),
                        idx=pgsql.Identifier(f"gist_{table}"),
                        tbl=pgsql.Identifier(table)))
        except Exception as e:
            gdf_log.warning(f'Could not (re)create gist index on '
                            f'{schema}.{table}(pos): {e}')
        model = self.TMD.SAParams.model()
        for item in (self.params_in_table or []):
            qIndex = self.TMD.SAParams.indexFromItem(item)
            model.removeRow(qIndex.row())
        self.params_in_list = 0

    def edit_tbl_name(self):
        """This function pops a question if the user wants to rename the selected tables to.
        Then it is replaced and the manual table is also updated."""
        for item in self.items_in_table:
            if item.checkState() == 2:
                schema, tbl = item.text().split('.')
                text, y_n = QInputDialog.getText(None, self.tr('Data set name'),
                                                 self.tr('What do you want to rename ') + tbl + self.tr(' to?'))
                if y_n:
                    query = pgsql.SQL("ALTER TABLE {schema}.{old} RENAME TO {new}").format(
                        schema=pgsql.Identifier(schema),
                        old=pgsql.Identifier(tbl),
                        new=pgsql.Identifier(text))
                    self.db.execute_sql(query)
                    query = pgsql.SQL("UPDATE {schema}.manual SET table_ = %s WHERE table_ = %s").format(
                        schema=pgsql.Identifier(schema))
                    self.db.execute_sql(query, params=(text, tbl))
        self.update_table_list()

    def edit_param_name(self):
        """Edit the name of all selected parameters."""
        for item in (self.params_in_table or []):
            if item.checkState() == 2:
                text, y_n = QInputDialog.getText(None, self.tr('Parameter name'),
                                                 self.tr('What do you want to rename ') + item.text() +
                                                 self.tr(' to?'))
                if y_n:
                    query = pgsql.SQL("ALTER TABLE {schema}.{tbl} RENAME {old_col} TO {new_col}").format(
                        schema=pgsql.Identifier(self.current_schema),
                        tbl=pgsql.Identifier(self.current_table),
                        old_col=pgsql.Identifier(item.text()),
                        new_col=pgsql.Identifier(text))
                    self.db.execute_sql(query)
                    try:
                        query = pgsql.SQL("ALTER INDEX {schema}.{old_idx} RENAME TO {new_idx}").format(
                            schema=pgsql.Identifier(self.current_schema),
                            old_idx=pgsql.Identifier(f"{item.text()}_{self.current_schema}_{self.current_table}"),
                            new_idx=pgsql.Identifier(f"{text}_{self.current_schema}_{self.current_table}"))
                        self.db.execute_sql(query)
                    except Exception as e:
                        gdf_log.warning(
                            f'Could not rename index for {item.text()}: {e}')
        self.retrieve_params()

    def update_table_list(self: Self) -> None:
        """Update the list of tables in the docket widget"""
        lw_list = self.parent.populate.get_lw_list()
        if self.tables_in_db != 0:
            model = self.TMD.SATables.model()
            for item in self.items_in_table:
                qIndex = self.TMD.SATables.indexFromItem(item)
                model.removeRow(qIndex.row())
        self.tables_in_db = 0
        for lw, schema in lw_list:
            table_names = self.db.get_tables_in_db(schema)
            for name in table_names:
                if str(name) in ['harrowing_manual', 'plowing_manual', 'manual']:
                    continue
                item_name = schema + '.' + str(name)
                testcase_name = QCoreApplication.translate("qadashboard", item_name, None)
                item = QListWidgetItem(testcase_name, self.TMD.SATables)
                item.setFlags(item.flags() | _item_flag('ItemIsUserCheckable'))
                item.setCheckState(_check_state('Unchecked'))
                self.tables_in_db += 1
        self.items_in_table = self.TMD.SATables.findItems('', _match_flag('MatchContains'))

    def remove_table_from_db(self: Self) -> None:
        """Removes the selected tables from the database"""
        msgBox = QMessageBox()
        msgBox.setText(self.tr('Do you really want to remove the selected tables from the database?'))
        msgBox.addButton(QPushButton(self.tr('Yes')), QMessageBox.ButtonRole.YesRole)
        msgBox.addButton(QPushButton(self.tr('No')), QMessageBox.ButtonRole.NoRole)
        if self.parent.test_mode:
            ret = 0
        else:
            ret = msgBox.exec()
        if ret == 1:
            return
        model = self.TMD.SATables.model()
        for item in self.items_in_table:
            if item.checkState() == 2:
                self.db.remove_table(item.text())
                qIndex = self.TMD.SATables.indexFromItem(item)
                model.removeRow(qIndex.row())
                self.tables_in_db -= 1
        self.items_in_table = self.TMD.SATables.findItems('', _match_flag('MatchContains'))

    def make_rows(self: Self) -> None:
        suc, s_table = self.check_multiple()
        if not suc:
            return
        schema, tbl = s_table.split('.')
        if schema == 'harvest':
            QMessageBox(None, self.tr('Error'),
                        self.tr('This option is not possible for harvest tables'))
            return
        query = pgsql.SQL(
            "SELECT field_row_id, row, course FROM {schema}.{tbl} ORDER BY row, field_row_id"
        ).format(schema=pgsql.Identifier(schema), tbl=pgsql.Identifier(tbl))
        row_courses = self.db.execute_and_return(query, return_failure=False)
        nr_rows = int(self.TMD.SBNumberOfRows.value())
        max_dev = int(self.TMD.SBMaxAngleOffset.value())
        avg_dist = float(self.TMD.SBAvgDistance.value())
        distance_bet_rows = float(self.TMD.SBRowDistance_2.value()) / 2
        stops = {}
        for i in range(1, nr_rows + 1):
            stops[i] = []

        current_course = 0
        for id_, row, course in row_courses:
            if len(stops[row]) == 0:
                stops[row].append(id_)
                current_course = course
                continue
            if round(course/max_dev) != round(current_course/max_dev):
                stops[row].append(id_)
                current_course = course
        for row in stops.keys():
            stops[row].append(len(row_courses)+1)
        task1 = QgsTask.fromFunction('Updating rows', make_rows, self.db, schema, tbl, avg_dist, distance_bet_rows, stops)
        self.parent.tsk_mngr.addTask(task1)
        

    def split_rows(self):
        """Spilt the harvest data to multiple rows."""
        suc, s_table = self.check_multiple()
        if not suc:
            return
        split_yield = self.TMD.CBSplitYield.isChecked()
        if split_yield:
            yield_row = self.TMD.CBColumns.currentText()
            if yield_row == self.tr('--- Select yield column ---'):
                QMessageBox(None, self.tr('Error'),
                            self.tr('In order to split the yield you need to specify the yield column'))
                return

        schema, tbl = s_table.split('.')
        if schema != 'harvest':
            QMessageBox(None, self.tr('Error'),
                        self.tr('This option is only possible for harvest tables'))
            return
        nbr_rows = int(self.TMD.CBNbrRows.currentText())
        spacing = float(self.TMD.SBRowDistance.text().replace(',', '.'))
        if nbr_rows == 2:
            move_d = [spacing / 2, spacing]
        elif nbr_rows == 4:
            move_d = [spacing + spacing / 2, spacing, spacing * 2, spacing * 3]
        else:  # nbr_rows == 6
            move_d = [2 * spacing + spacing / 2, spacing, spacing * 2,
                      spacing * 3, spacing * 4, spacing * 5]
        tbl_id = pgsql.Identifier(tbl)
        min_row_id = self.db.execute_and_return(
            pgsql.SQL("SELECT min(field_row_id) FROM harvest.{tbl}").format(tbl=tbl_id))[0][0]
        max_row_id = self.db.execute_and_return(
            pgsql.SQL("SELECT max(field_row_id) FROM harvest.{tbl}").format(tbl=tbl_id))[0][0]
        org_min = min_row_id
        org_max = max_row_id
        if min_row_id is None:
            # If the user choose the wrong field and this part cant be used.
            return [False, 'Wrong field']
        for row_nbr in range(nbr_rows):
            if row_nbr == 0:
                bearing = 270
            else:
                min_row_id, max_row_id = self.duplicate_first_row(org_min, org_max, tbl)
                bearing = 90
            distance = move_d[row_nbr]
            harvest_tbl = pgsql.SQL("harvest.{}").format(tbl_id)
            for i in range(min_row_id, max_row_id, 2000):
                query = pgsql.SQL(
                    "WITH first_selected AS ("
                    " SELECT field_row_id, pos, st_azimuth(pos,"
                    " (SELECT pos FROM {tbl} new WHERE org.field_row_id+1 = new.field_row_id)) AS bearing"
                    " FROM {tbl} org"
                    " WHERE org.field_row_id >= %s - 3 AND org.field_row_id < %s + 3),"
                    " sub_section AS ("
                    " SELECT field_row_id,"
                    " st_project(pos::geography, %s,"
                    " (SELECT atan2(avg(sin(bearing)), avg(cos(bearing)))"
                    " FROM first_selected fs"
                    " WHERE fs.field_row_id > (ss.field_row_id - 3)"
                    " AND fs.field_row_id < (ss.field_row_id + 3)) + radians(%s))::geometry AS new_pos"
                    " FROM first_selected ss)"
                    " UPDATE {tbl} SET pos = new_pos"
                    " FROM sub_section"
                    " WHERE {tbl}.field_row_id = sub_section.field_row_id"
                    " AND {tbl}.field_row_id >= %s"
                    " AND {tbl}.field_row_id < %s"
                ).format(tbl=harvest_tbl)
                self.db.execute_sql(query, params=(i, i + 2000, distance, bearing, i, i + 2000))
            if split_yield:
                query = pgsql.SQL("UPDATE {schema}.{tbl} SET {col} = {col} / %s").format(
                    schema=pgsql.Identifier(schema),
                    tbl=pgsql.Identifier(tbl),
                    col=pgsql.Identifier(yield_row))
                self.db.execute_sql(query, params=(nbr_rows,))

    def update_column_list(self):
        suc, s_table = self.check_multiple()
        if not suc:
            return
        schema, tbl = s_table.split('.')
        columns = self.db.get_all_columns(tbl, schema, "'field_row_id'")
        lw = self.TMD.CBColumns
        lw.clear()
        lw.addItem(self.tr('--- Select yield column ---'))
        for name in columns:
            lw.addItem(str(name))

    def duplicate_first_row(self, org_min, org_max, table, schema='harvest'):
        """Duplicates the original data table and returns the first and last row_id"""
        tbl_id = pgsql.SQL("{}.{}").format(
            pgsql.Identifier(schema), pgsql.Identifier(table))
        sql_max_row = pgsql.SQL("SELECT max(field_row_id) FROM {tbl}").format(tbl=tbl_id)
        current_max = self.db.execute_and_return(sql_max_row)[0][0]
        columns = self.db.get_all_columns(table, schema)
        c1_parts = pgsql.SQL(", ").join(pgsql.Identifier(c) for c in columns)
        c2_parts = pgsql.SQL(", ").join(
            pgsql.SQL("ROW_NUMBER() OVER () + %s") if c == 'field_row_id'
            else pgsql.Identifier(c)
            for c in columns)
        offset = current_max + 1
        c2_params = [offset] if 'field_row_id' in columns else []
        query = pgsql.SQL(
            "INSERT INTO {tbl} ({c1}) SELECT {c2} FROM {tbl}"
            " WHERE field_row_id >= %s AND field_row_id <= %s"
        ).format(tbl=tbl_id, c1=c1_parts, c2=c2_parts)
        self.db.execute_sql(query, params=tuple(c2_params + [org_min, org_max]))
        now_max = self.db.execute_and_return(sql_max_row)[0][0]
        return current_max + 1, now_max


def make_rows(db, schema, tbl, avg_dist, distance_bet_rows, stops):
        tbl_id = pgsql.SQL("{}.{}").format(
            pgsql.Identifier(schema), pgsql.Identifier(tbl))
        for j, (row, stops_) in enumerate(stops.items()):
            for i, stop in enumerate(stops_[:-1]):
                query = pgsql.SQL(
                    "WITH sel AS ("
                    " SELECT st_buffer(st_makeline(pos ORDER BY field_row_id)::geography, %s) AS outer_row"
                    " FROM {tbl}"
                    " WHERE row = %s AND field_row_id < %s AND field_row_id >= %s)"
                    " UPDATE {tbl}"
                    " SET polygon = st_multi(st_intersection("
                    " st_buffer(pos::geography, %s), outer_row)::geometry)"
                    " FROM sel"
                    " WHERE row = %s AND field_row_id < %s AND field_row_id >= %s"
                ).format(tbl=tbl_id)
                db.execute_sql(query, params=(
                    distance_bet_rows, row, stops[row][i+1], stop,
                    avg_dist, row, stops[row][i+1], stop))
