from qgis.PyQt.QtWidgets import (
    QGridLayout,
    QMessageBox,
    QTableView,
    QWidget,
    QWizardPage,
)

from openlog.datamodel.assay.categories import CategoriesTableDefinition
from openlog.datamodel.assay.generic_assay import AssayDatabaseDefinition
from openlog.datamodel.connection.openlog_connection import OpenLogConnection
from openlog.gui.import_assay.assay_database_definition_widget import (
    AssayDatabaseDefinitionWidget,
)
from openlog.gui.import_assay.wizard.assay_import import AssayImportPageWizard
from openlog.gui.import_assay.wizard.assay_selection import AssaySelectionPageWizard
from openlog.toolbelt import PlgTranslator
from openlog.toolbelt.string_map_table_model import StringMapTableModel


class AssayDatabaseDefinitionPageWizard(QWizardPage):
    def __init__(
        self,
        parent: QWidget,
        openlog_connection: OpenLogConnection,
        assay_selection_wizard: AssaySelectionPageWizard,
        assay_import_wizard: AssayImportPageWizard,
    ) -> None:
        """
        QWizard to define assay database definition

        Args:
            openlog_connection: OpenLogConnection
            assay_selection_wizard : AssaySelectionPageWizard defining assay
            assay_import_wizard : AssayImportPageWizard defining assay data
            parent : QWidget parent
        """
        super().__init__(parent)
        self.tr = PlgTranslator().tr
        self._openlog_connection = openlog_connection
        self._assay_selection_wizard = assay_selection_wizard
        self._assay_import_wizard = assay_import_wizard
        self.setTitle(self.tr("Downhole data database definition"))

        layout = QGridLayout()
        self.assay_database_definition_widget = AssayDatabaseDefinitionWidget()
        layout.addWidget(self.assay_database_definition_widget, 0, 0)

        self.category_definition_table = QTableView(self)
        self.category_definition_table.horizontalHeader().setVisible(False)
        self.category_definition_table_model = StringMapTableModel(
            self, self.tr("Category"), self.tr("Table Name")
        )
        self.category_definition_table.setModel(self.category_definition_table_model)
        layout.addWidget(self.category_definition_table, 1, 0)

        self.setLayout(layout)

    def initializePage(self) -> None:
        """
        Initialize page before show.

        Update AssayDatabaseDefinitionWidget with data from assay selection

        """

        self.assay_database_definition_widget.setEnabled(
            self._assay_selection_wizard.assay_creation_needed()
        )

        assay_definition = self._assay_import_wizard.get_assay_definition()

        # Define default schema and table name
        database_definition = (
            self.assay_database_definition_widget.get_assay_database_definition()
        )
        database_definition.schema = (
            self._openlog_connection.get_assay_iface().default_assay_schema()
        )
        database_definition.table_name = assay_definition.variable
        if not self._assay_selection_wizard.assay_creation_needed():
            database_definition = (
                self._openlog_connection.get_assay_iface().get_assay_table_definition(
                    assay_definition.variable
                )
            )
        else:
            current_cols = assay_definition.columns.keys()
            database_definition_cols = database_definition.y_col.keys()

            # Add default val for undefined columns
            undefined_columns = [
                col for col in current_cols if col not in database_definition_cols
            ]
            for col in undefined_columns:
                database_definition.y_col[col] = col

            # Removed unused columns
            unused_cols = [
                col for col in database_definition_cols if col not in current_cols
            ]
            for col in unused_cols:
                del database_definition.y_col[col]

        self.assay_database_definition_widget.set_assay_database_definition(
            database_definition
        )

        # Get categories to be created
        current_categories = [
            c.name
            for c in self._openlog_connection.get_categories_iface().get_available_categories_table()
        ]

        created_category = [
            c
            for c in self._assay_import_wizard.dataset_edit.get_used_category()
            if c not in current_categories
        ]
        cat_map = {}
        for c in created_category:
            cat_map[c] = c.replace(" ", "_")
        self.category_definition_table_model.set_string_map(cat_map)

    def validatePage(self) -> bool:
        """
        Validate page by checking that a valid assay database definition is used

        Returns: True if a valid assay database definition is used, False otherwise

        """
        valid = True

        if (
            not self.assay_database_definition_widget.get_assay_database_definition().is_valid()
        ):
            valid = False
            QMessageBox.warning(
                self,
                self.tr("Invalid downhole data database definition"),
                self.tr(f"Check for blank characters in downhole data database definition."),
            )

        if valid:
            for (
                category,
                category_table,
            ) in self.category_definition_table_model.get_string_map().items():
                valid &= " " not in category_table and category_table != ""
            if not valid:
                QMessageBox.warning(
                    self,
                    self.tr("Invalid category table name"),
                    self.tr(f"Check for space in category table name."),
                )

        return valid

    def get_assay_database_definition(self) -> AssayDatabaseDefinition:
        """
        Returns AssayDatabaseDefinition defined in wizard page

        Returns: AssayDatabaseDefinition

        """
        return self.assay_database_definition_widget.get_assay_database_definition()

    def get_created_categories(self) -> [CategoriesTableDefinition]:
        # Get values from model
        cat_map = self.category_definition_table_model.get_string_map()
        res = []
        for category, category_table in cat_map.items():
            res.append(
                CategoriesTableDefinition(name=category, table_name=category_table)
            )
        return res
