from qgis.PyQt.QtWidgets import (
    QHeaderView,
    QLabel,
    QMessageBox,
    QTreeWidget,
    QTreeWidgetItem,
    QVBoxLayout,
    QWidget,
    QWizard,
    QWizardPage,
)

from openlog.datamodel.connection.interfaces.assay_interface import AssayInterface
from openlog.datamodel.connection.openlog_connection import OpenLogConnection
from openlog.datamodel.imports.import_classes import ImportType
from openlog.gui.import_assay.wizard.assay_database_definition import (
    AssayDatabaseDefinitionPageWizard,
)
from openlog.gui.import_assay.wizard.assay_import import AssayImportPageWizard
from openlog.gui.import_assay.wizard.assay_selection import AssaySelectionPageWizard
from openlog.gui.import_data.wizard.person_dataset import (
    PersonDatasetSelectionPageWizard,
)
from openlog.plugins.manager import get_plugin_manager


class ConfirmationPageWizard(QWizardPage):
    def __init__(
        self,
        parent: QWidget,
        openlog_connection: OpenLogConnection,
        assay_selection_wizard: AssaySelectionPageWizard,
        assay_database_definition_wizard: AssayDatabaseDefinitionPageWizard,
        assay_import_wizard: AssayImportPageWizard,
    ) -> None:
        """
        QWizardPage used to confirm assay imports and commit database session

        Args:
            parent : QWidget parent
            openlog_connection: OpenLogConnection used to import data
            assay_selection_wizard: AssaySelectionPageWizard
            assay_database_definition_wizard: AssayDatabaseDefinitionPageWizard
            assay_import_wizard: AssayImportPageWizard
        """
        super().__init__(parent)
        self._openlog_connection = openlog_connection
        self._assay_selection_wizard = assay_selection_wizard
        self._assay_database_definition_wizard = assay_database_definition_wizard
        self._assay_import_wizard = assay_import_wizard
        self.skip_validation = False
        self.setTitle("Import summary")

        label = QLabel(
            "Here are the data that will be imported into openlog connection."
        )
        label.setWordWrap(True)

        self.layout = QVBoxLayout()
        self.layout.addWidget(label)

        self.tree_widget = QTreeWidget()
        self.tree_widget.setHeaderHidden(True)
        self.tree_widget.header().setSectionResizeMode(
            QHeaderView.ResizeMode.ResizeToContents
        )
        self.layout.addWidget(self.tree_widget)

        self.setLayout(self.layout)

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

        Update tree widget with imported data

        """
        self.tree_widget.clear()
        self.tree_widget.setColumnCount(2)

        self._add_assay_definition_tree_item()
        self._add_assay_database_definition()
        self._add_categories_definition()
        self._add_assay_import_tree_item()

        self.tree_widget.resizeColumnToContents(0)
        self.tree_widget.resizeColumnToContents(1)

    def _add_assay_definition_tree_item(self) -> None:
        """
        Add QTreeWidgetItem for assay definition

        """
        assay_definition = self._assay_import_wizard.get_assay_definition()
        assay_definition_tw = QTreeWidgetItem(self.tree_widget)
        assay_definition_tw.setExpanded(True)
        assay_definition_tw.setText(0, "Downhole data")

        self._add_tree_item(assay_definition_tw, "Variable", assay_definition.variable)
        self._add_tree_item(
            assay_definition_tw, "Domain", assay_definition.domain.value
        )
        self._add_tree_item(
            assay_definition_tw, "Extent", assay_definition.data_extent.value
        )

        column_tw = self._add_tree_item(
            assay_definition_tw, "Columns", str(len(assay_definition.columns))
        )
        for col, val in assay_definition.columns.items():
            self._add_tree_item(column_tw, col, val.series_type.value)

    def _add_categories_definition(self) -> None:
        created_categories = (
            self._assay_database_definition_wizard.get_created_categories()
        )
        if len(created_categories):
            category_tw = QTreeWidgetItem(self.tree_widget)
            category_tw.setExpanded(True)
            category_tw.setText(0, "Category")
            for category in created_categories:
                cat_tw = self._add_tree_item(category_tw, category.name, "")
                self._add_tree_item(cat_tw, "Table name", category.table_name)
                self._add_tree_item(cat_tw, "Name column", category.name_column)
                self._add_tree_item(cat_tw, "Schema", category.schema)

    def _add_assay_database_definition(self) -> None:
        """
        Add QTreeWidgetItem for assay database definition

        """
        assay_database_definition = (
            self._assay_database_definition_wizard.get_assay_database_definition()
        )
        assay_database_definition_tw = QTreeWidgetItem(self.tree_widget)
        assay_database_definition_tw.setExpanded(True)
        assay_database_definition_tw.setText(0, "Database definition")

        self._add_tree_item(
            assay_database_definition_tw,
            "Table",
            assay_database_definition.table_name,
        )
        self._add_tree_item(
            assay_database_definition_tw,
            "Hole ID column",
            assay_database_definition.hole_id_col,
        )
        self._add_tree_item(
            assay_database_definition_tw,
            "Dataset column",
            assay_database_definition.dataset_col,
        )
        self._add_tree_item(
            assay_database_definition_tw,
            "X column",
            assay_database_definition.x_col,
        )

        column_tw = self._add_tree_item(
            assay_database_definition_tw,
            "Columns",
            str(len(assay_database_definition.y_col)),
        )
        for col, val in assay_database_definition.y_col.items():
            self._add_tree_item(column_tw, col, val)

        self._add_tree_item(
            assay_database_definition_tw,
            "X end column",
            assay_database_definition.x_end_col,
        )
        self._add_tree_item(
            assay_database_definition_tw,
            "Schema",
            assay_database_definition.schema,
        )

    def _add_assay_import_tree_item(self) -> None:
        """
        Add QTreeWidgetItem for assay imported data

        """
        assay_import_tw = QTreeWidgetItem(self.tree_widget)
        assay_import_tw.setExpanded(True)
        assay_import_tw.setText(0, "Imported data")

        assays = self._assay_import_wizard.get_assays()
        self._add_tree_item(assay_import_tw, "Collar count", str(len(assays)))
        data_count = 0
        for assay in assays:
            data_count += len(assay.import_x_values)
        self._add_tree_item(assay_import_tw, "Data count", str(data_count))

        if self._assay_import_wizard.dataset_edit.is_extended_assay:
            self._add_tree_item(
                assay_import_tw,
                "Removed rows (thickness of 0)",
                str(self._assay_import_wizard.dataset_edit.removed_rows),
            )

    @staticmethod
    def _add_tree_item(
        parent: QTreeWidgetItem, name: str, value: str
    ) -> QTreeWidgetItem:
        """
        Create a QTreeWidget item with name and value column

        Args:
            parent: parent QTreeWidgetItem
            name: name
            value: value

        Returns: created QTreeWidgetItem

        """
        item = QTreeWidgetItem(parent)
        item.setText(0, name)
        item.setText(1, value)
        return item

    def validatePage(self) -> bool:
        """
        Import data

        Returns: True if no import error occured, False otherwise

        """
        if self.skip_validation:
            return super().validatePage()
        valid = True
        try:
            categories_iface = self._openlog_connection.get_categories_iface()
            if self._assay_selection_wizard.assay_creation_needed():
                # Add new assay definition and database definition
                assay_definition = self._assay_import_wizard.get_assay_definition()
                assay_database_definition = (
                    self._assay_database_definition_wizard.get_assay_database_definition()
                )

                self._openlog_connection.get_assay_iface().add_assay_table(
                    assay_definition, assay_database_definition
                )

                # Add new categories
                created_categories = (
                    self._assay_database_definition_wizard.get_created_categories()
                )
                categories_iface.import_categories_table(created_categories)

            # Add used categories values
            cat_values = self._assay_import_wizard.get_used_categories_values()
            for table_name, value in cat_values.items():
                cat = categories_iface.get_categories_table(table_name)
                categories_iface.import_categories_data(cat, value)

            self._openlog_connection.get_assay_iface().import_assay_data(
                self._assay_import_wizard.get_assays(),
                self._assay_import_wizard.get_import_policies(),
            )

            self._openlog_connection.commit()

        except AssayInterface.ImportException as exc:
            valid = False
            QMessageBox.warning(self, "Import exception", str(exc))
            self._openlog_connection.rollback()
        return valid


class AssayImportWizard(QWizard):
    def rejected(self) -> None:
        """
        Rollback openlog connection if data import is rejected

        """
        self.openlog_connection.rollback()
        super(AssayImportWizard).rejected()

    def __init__(
        self, openlog_connection: OpenLogConnection, parent: QWidget = None
    ) -> None:
        """
        QWizard to import assay into openlog connection from .csv files.

        Args:
            openlog_connection: OpenLogConnection used to import data
            parent : QWidget parent
        """
        super().__init__(parent)
        self.openlog_connection = openlog_connection
        self.wizard_type = ImportType.ASSAY
        # import profile actions
        self.profile_generator = (
            get_plugin_manager().get_advanced_imports_plugin().profile_generator(self)
        )
        self.profile_enabled = get_plugin_manager().get_advanced_imports_plugin().enable
        self.profile_btn = self.profile_generator.get_wizard_button()
        self.profile_btn.setEnabled(False)
        if not self.profile_enabled:
            self.profile_btn.setToolTip("Premium only")

        self.setWindowTitle("Downhole data import")

        # person and dataset
        self.person_dataset_import_wizard = PersonDatasetSelectionPageWizard(
            self, self.openlog_connection
        )
        self.addPage(self.person_dataset_import_wizard)

        self.assay_selection_wizard = AssaySelectionPageWizard(
            self, self.openlog_connection
        )
        self.addPage(self.assay_selection_wizard)

        self.assay_import_wizard = AssayImportPageWizard(
            self,
            self.openlog_connection,
            self.assay_selection_wizard,
            self.person_dataset_import_wizard,
        )
        self.addPage(self.assay_import_wizard)

        self.assay_selection_wizard.assay_definition_widget.data_extent_cb.currentIndexChanged.connect(
            self.assay_import_wizard._set_import_options
        )

        self.assay_database_definition_wizard = AssayDatabaseDefinitionPageWizard(
            self,
            self.openlog_connection,
            self.assay_selection_wizard,
            self.assay_import_wizard,
        )
        self.addPage(self.assay_database_definition_wizard)

        self.confirmation_page = ConfirmationPageWizard(
            self,
            self.openlog_connection,
            self.assay_selection_wizard,
            self.assay_database_definition_wizard,
            self.assay_import_wizard,
        )

        # Wizard to confirm data import and commit database connection
        self.addPage(self.confirmation_page)

        # profile button appears on import_wizard page
        self.currentIdChanged.connect(self._update_buttons)
        # profile button is enabled if data source is set
        self.assay_import_wizard.completeChanged.connect(self._update_buttons)

    def _update_buttons(self) -> None:
        """
        Update wizard buttons.
        """
        self.setOption(QWizard.HaveHelpButton, False)
        self.profile_btn.setEnabled(False)

        if isinstance(self.currentPage(), AssayImportPageWizard):
            self.setOption(QWizard.HaveHelpButton, True)
            self.setButton(QWizard.HelpButton, self.profile_btn)
            # appears only if source data is set + already created assay
            if (
                self.profile_enabled
                and self.assay_selection_wizard.assay_selection_rb.isChecked()
            ):
                self.profile_btn.setEnabled(True)
