from qgis.PyQt import QtCore
from qgis.PyQt.QtCore import Qt
from qgis.PyQt.QtGui import QStandardItemModel

from openlog.gui.assay_visualization.stacked.stacked_config import StackedConfiguration
from openlog.gui.assay_visualization.stacked.stacked_utils import invalid_series_type


class StackConfigTableModel(QStandardItemModel):
    HOLE_ID_COL = 0
    ASSAY_COL = 1
    ASSAY_COLUMN_COL = 2
    UNIT_COL = 3
    SERIES_TYPE_COL = 4

    CONFIG_ROLE = Qt.UserRole

    def __init__(self, parent=None) -> None:
        """
        QStandardItemModel for StackedConfiguration display

        Args:
            parent: QWidget
            openlog_connection: OpenLogConnection used to get collar display name
        """
        super().__init__(parent)
        self.setHorizontalHeaderLabels(
            [
                "Hole ID",
                "Downhole data",
                "Column",
                "Unit",
                "Type",
            ]
        )

    def flags(self, index: QtCore.QModelIndex) -> QtCore.Qt.ItemFlags:
        """
        Override QStandardItemModel flags.

        Args:
            index: QModelIndex

        Returns: index flags

        """
        flags = super().flags(index)
        flags = flags | QtCore.Qt.ItemIsUserCheckable
        return flags

    def set_stacked_config(self, stacked_config: StackedConfiguration):
        for i in reversed(range(self.rowCount())):
            self.removeRow(i)

        for config in stacked_config.config_list:
            self.insertRow(self.rowCount())
            row = self.rowCount() - 1
            self.setData(self.index(row, self.HOLE_ID_COL), config.hole_display_name)
            self.setData(self.index(row, self.HOLE_ID_COL), config, self.CONFIG_ROLE)
            # By default invalid series type are unchecked
            check_role = (
                Qt.Unchecked
                if config.column.series_type in invalid_series_type()
                else Qt.Checked
            )
            self.setData(
                self.index(row, self.HOLE_ID_COL), check_role, Qt.CheckStateRole
            )
            self.setData(self.index(row, self.ASSAY_COL), config.assay_display_name)
            self.setData(self.index(row, self.ASSAY_COLUMN_COL), config.column.name)
            self.setData(self.index(row, self.UNIT_COL), str(config.column.unit))
            self.setData(
                self.index(row, self.SERIES_TYPE_COL), config.column.series_type.name
            )

    def get_stacked_config(self, only_checked: bool = True) -> StackedConfiguration:
        stacked_config = StackedConfiguration()

        for i in range(self.rowCount()):
            config = self.data(self.index(i, self.HOLE_ID_COL), self.CONFIG_ROLE)
            if only_checked:
                if (
                    self.data(self.index(i, self.HOLE_ID_COL), Qt.CheckStateRole)
                    == Qt.Checked
                ):
                    stacked_config.add_column_config(config)
            else:
                stacked_config.add_column_config(config)

        return stacked_config
