from qgis.PyQt import QtCore
from qgis.PyQt.QtCore import QVariant
from qgis.PyQt.QtGui import QColor, QFont, QStandardItemModel

from openlog.datamodel.assay.generic_assay import AssaySeriesType
from openlog.datamodel.assay.spherical_assay import SphericalDefinition
from openlog.gui.utils.column_definition import ColumnDefinition


class SphericalDataTableModel(QStandardItemModel):
    COLUMN_COL = 0
    TYPE_COL = 1
    DIP_COL = 2
    AZIMUTH_COL = 3
    POLARITY_COL = 4

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

        Args:
            parent: QWidget
        """
        super().__init__(parent)
        self.setHorizontalHeaderLabels(
            [
                self.tr("Column"),
                self.tr("Type"),
                self.tr("Dip"),
                self.tr("Azimuth"),
                self.tr("Polarity"),
            ]
        )

    def flags(self, index: QtCore.QModelIndex) -> QtCore.Qt.ItemFlags:
        """
        Override QStandardItemModel flags to remove edition for column col

        Args:
            index: QModelIndex

        Returns: index flags

        """
        flags = super().flags(index)
        if index.column() == self.COLUMN_COL:
            flags = flags & QtCore.Qt.ItemIsEditable
        return flags

    def data(
        self, index: QtCore.QModelIndex, role: int = QtCore.Qt.DisplayRole
    ) -> QVariant:
        """
        Override QStandardItemModel data() for font of expected column

        Args:
            index: QModelIndex
            role: Qt role

        Returns: QVariant

        """
        result = super().data(index, role)
        if role == QtCore.Qt.FontRole and index.column() == self.COLUMN_COL:
            result = QFont()
            result.setBold(True)

        return result

    def get_spherical_assay(self) -> dict[str:SphericalDefinition]:

        result = {}
        for i in range(self.rowCount()):
            expected_col = str(self.data(self.index(i, self.COLUMN_COL)))
            if expected_col:
                dip_column = self.data(self.index(i, self.DIP_COL))
                azimuth_col = self.data(self.index(i, self.AZIMUTH_COL))
                polarity_col = self.data(self.index(i, self.POLARITY_COL))

                spherical_type = self.data(self.index(i, self.TYPE_COL))

                column = SphericalDefinition(
                    dip_col=dip_column,
                    azimuth_col=azimuth_col,
                    polarity_col=polarity_col,
                    type_=spherical_type,
                )
                result[expected_col] = column
        return result

    def set_spherical_assay(self, spherical: dict[str:SphericalDefinition]) -> None:
        """
        Define expected column and remove not used column

        Args:
            column_uncertainty:
        """
        columns_name = spherical.keys()
        for i in reversed(range(self.rowCount())):
            expected_col = str(self.data(self.index(i, self.COLUMN_COL)))
            if expected_col not in columns_name:
                self.removeRow(i)

        for col, uncertainty in spherical.items():
            row_index = self._get_column_row_index(col)
            if row_index != -1:
                self._set_spherical_assay(row_index, col, uncertainty)
            else:
                self.add_spherical_assay(col, uncertainty)

    def remove_column(self, column: str) -> None:
        """
        Remove a column spherical assay definition

        Args:
            column: (str) column name

        """
        for i in range(self.rowCount()):
            column_row = str(self.data(self.index(i, self.COLUMN_COL)))
            if column_row == column:
                self.removeRow(i)
                return

    def _get_column_row_index(self, col: str) -> int:
        """
        Get expected column column index (-1 if expected column not available)

        Args:
            col: (str) expected column

        Returns: expected column column index (-1 if expected column not available)

        """
        row = -1
        for i in range(self.rowCount()):
            val = str(self.data(self.index(i, self.COLUMN_COL)))
            if val == col:
                row = i
        return row

    def _get_or_create_column(self, col: str) -> int:
        """
        Get or create column

        Args:
            col: expected column name

        Returns: column index

        """
        row = self._get_column_row_index(col)
        if row == -1:
            self.add_spherical_assay(column=col, spherical=SphericalDefinition())
            row = self.rowCount() - 1
        return row

    def _set_spherical_assay(
        self, row: int, column: str, spherical: SphericalDefinition
    ):

        self.setData(self.index(row, self.COLUMN_COL), column)

        self.setData(
            self.index(row, self.DIP_COL),
            spherical.dip_col,
        )
        self.setData(
            self.index(row, self.AZIMUTH_COL),
            spherical.azimuth_col,
        )
        self.setData(self.index(row, self.TYPE_COL), spherical.type_)

        # set item non editable if line
        polarity = spherical.polarity_col if spherical.type_ == "PLANE" else ""
        self.setData(self.index(row, self.POLARITY_COL), polarity)
        item = self.item(row, self.POLARITY_COL)

        if spherical.type_ == "LINE":
            item.setFlags(item.flags() & ~QtCore.Qt.ItemIsEditable)
        else:
            item.setFlags(item.flags() | QtCore.Qt.ItemIsEditable)

    def add_spherical_assay(self, column: str, spherical: SphericalDefinition):

        """
        Add a row in QStandardItemModel

        Args:
            column : (str) column name
            uncertainty: AssayColumnUncertainty
        """
        self.insertRow(self.rowCount())
        row = self.rowCount() - 1
        self._set_spherical_assay(row, column, spherical)
