from qgis.PyQt.QtCore import QObject
from qgis.PyQt.QtGui import QStandardItemModel
from xplordb.datamodel.collar import Collar
from xplordb.datamodel.survey import Survey


class SurveyCreationLayerModel(QStandardItemModel):
    DEPTH_COL = 0
    DIP_COL = 1
    AZIMUTH_COL = 2

    def __init__(self, parent: QObject = None):
        """
        QStandardItemModel for survey creation.

        Args:
            parent: QObject parent
        """
        super(SurveyCreationLayerModel, self).__init__(parent)
        self.setHorizontalHeaderLabels(
            [self.tr("Depth"), self.tr("Dip"), self.tr("Azimuth")]
        )

    def get_survey_list(self) -> [Survey]:
        """
        Get survey list (warning no person and dataset are defined)

        Returns: survey list (warning no person and dataset are defined)

        """
        result = []
        for i in range(0, self.rowCount()):
            result.append(self._survey(i))
        return result

    def add_survey(self, survey: Survey = None, collar: Collar = None):
        """
        Add row in model from Survey

        Args:
            survey: Survey
        """
        self.insertRow(self.rowCount())
        if survey:
            self._set_survey_data(self.rowCount() - 1, survey, collar)

    def _survey(self, row: int) -> Survey:
        """
        Get survey definition for specific row

        Args:
            row: model row

        Returns: Survey for specific row

        """
        survey = Survey(
            hole_id="",  # Hole ID must be defined by widget using model
            data_set="",  # Data set must be defined by widget using model
            loaded_by="",  # Loaded by person must be defined by widget using model
            dip=self.data(self.index(row, self.DIP_COL)),
            azimuth=self.data(self.index(row, self.AZIMUTH_COL)),
            depth=self.data(self.index(row, self.DEPTH_COL)),
        )
        return survey

    def _set_survey_data(self, row: int, survey: Survey, collar: Collar = None) -> None:
        """
        Define QStandardItemModel data from Survey

        Args:
            row: model row
            survey: Survey
        """
        self.setData(self.index(row, self.DIP_COL), survey.dip)
        self.setData(self.index(row, self.AZIMUTH_COL), survey.azimuth)
        if collar:
            self.setData(self.index(row, self.DEPTH_COL), collar.planned_eoh)
        else:
            self.setData(self.index(row, self.DEPTH_COL), survey.depth)
