from qgis.PyQt.QtWidgets import QLabel, QMessageBox, QVBoxLayout, QWidget, QWizardPage
from xplordb.datamodel.collar import Collar
from xplordb.datamodel.metadata import RawCollarMetadata

from openlog.datamodel.assay.generic_assay import AssaySeriesType
from openlog.datamodel.connection.openlog_connection import OpenLogConnection
from openlog.gui.utils.column_definition import ColumnDefinition
from openlog.gui.utils.delimited_text_import_widget import DelimitedTextImportWidget
from openlog.toolbelt.translator import PlgTranslator

BASE_SETTINGS_KEY = "/OpenLog/gui/import/collar"


class CollarsImportPageWizard(QWizardPage):
    def __init__(self, parent: QWidget, openlog_connection: OpenLogConnection) -> None:
        """
        QWizard to import collars into xplordb from csv file

        Args:
            openlog_connection: OpenLogConnection used to import collar
            parent : QWidget parent
        """
        super().__init__(parent)
        self.tr = PlgTranslator().tr
        self._openlog_connection = openlog_connection
        self.setTitle(self.tr("Collar import"))

        self.HOLE_ID_COL = self.tr("HoleID*")
        self.X_COL = self.tr("Easting*")
        self.Y_COL = self.tr("Northing*")
        self.Z_COL = self.tr("Elevation")
        self.EOH_COL = self.tr("EOH")
        self.PLANNED_EOH_COL = self.tr("Planned EOH")
        self.DIP = self.tr("Dip")
        self.AZIMUTH = self.tr("Azimuth")
        self.SURVEY_DATE_COL = self.tr("Date")

        label = QLabel(self.tr("Select a .csv file and define CRS for collars import"))
        label.setWordWrap(True)

        layout = QVBoxLayout()
        layout.addWidget(label)
        self.dataset_edit = DelimitedTextImportWidget()
        self.dataset_edit.enable_crs_selection(True)
        self.dataset_edit.enable_elevation_from_dtm(
            True, self.X_COL, self.Y_COL, self.Z_COL
        )
        self.dataset_edit.set_column_definition(
            [
                ColumnDefinition(
                    column=self.HOLE_ID_COL,
                    fixed=True,
                    series_type=AssaySeriesType.NOMINAL,
                ),
                ColumnDefinition(
                    column=self.X_COL,
                    unit="m",
                    fixed=True,
                    series_type=AssaySeriesType.NUMERICAL,
                ),
                ColumnDefinition(
                    column=self.Y_COL,
                    unit="m",
                    fixed=True,
                    series_type=AssaySeriesType.NUMERICAL,
                ),
                ColumnDefinition(
                    column=self.Z_COL,
                    unit="m",
                    fixed=True,
                    optional=True,
                    series_type=AssaySeriesType.NUMERICAL,
                ),
                ColumnDefinition(
                    column=self.EOH_COL,
                    unit="m",
                    fixed=True,
                    optional=True,
                    series_type=AssaySeriesType.NUMERICAL,
                ),
                ColumnDefinition(
                    column=self.PLANNED_EOH_COL,
                    unit="m",
                    fixed=True,
                    optional=True,
                    series_type=AssaySeriesType.NUMERICAL,
                ),
                ColumnDefinition(
                    column=self.DIP,
                    unit="°",
                    fixed=True,
                    optional=True,
                    series_type=AssaySeriesType.NUMERICAL,
                ),
                ColumnDefinition(
                    column=self.AZIMUTH,
                    unit="°",
                    fixed=True,
                    optional=True,
                    series_type=AssaySeriesType.NUMERICAL,
                ),
                ColumnDefinition(
                    column=self.SURVEY_DATE_COL,
                    unit="datetime",
                    fixed=True,
                    optional=True,
                    series_type=AssaySeriesType.DATETIME,
                ),
            ]
        )
        layout.addWidget(self.dataset_edit)
        self.setLayout(layout)
        self.dataset_edit.restore_settings(BASE_SETTINGS_KEY)

    def data_label(self) -> str:
        """
        Returns label to be used in confirmation dialog

        Returns: imported data label

        """
        return self.tr("Collars")

    def data_count(self) -> int:
        """
        Returns expected imported data count to be displayed in confirmation dialog

        Returns: expected imported data count

        """
        df = self.dataset_edit.get_dataframe()
        return df.shape[0] if df is not None else 0

    def import_data(self):
        """
        Import data into openlog database.

        OpenLogConnection.ImportData exception can be raised.

        """
        df = self.dataset_edit.get_dataframe()
        if df is not None:
            collars = [
                Collar(
                    hole_id=r[self.HOLE_ID_COL],
                    data_set=self.field("dataset"),
                    loaded_by=self.field("person"),
                    x=r[self.X_COL],
                    y=r[self.Y_COL],
                    z=r[self.Z_COL] if r[self.Z_COL] is not None else 0.0,
                    srid=self.dataset_edit.crs().postgisSrid(),
                    project_srid=self._openlog_connection.get_layers_iface()
                    .get_collar_layer()
                    .crs()
                    .postgisSrid(),
                    eoh=r[self.EOH_COL],
                    planned_eoh=r[self.PLANNED_EOH_COL],
                    dip=r[self.DIP] if r[self.DIP] else -90.0,
                    azimuth=r[self.AZIMUTH] if r[self.AZIMUTH] else 0.0,
                    survey_date=r[self.SURVEY_DATE_COL],
                )
                for index, r in df.iterrows()
            ]

            # find extra columns
            default_cols = [
                self.HOLE_ID_COL,
                self.X_COL,
                self.Y_COL,
                self.Z_COL,
                self.EOH_COL,
                self.PLANNED_EOH_COL,
                self.DIP,
                self.AZIMUTH,
                self.SURVEY_DATE_COL,
            ]
            extra_cols = [col for col in df.columns if col not in default_cols]

            # set values into RawCollarMetadata instances
            metadatas = []
            for collar, (_, values_serie) in zip(collars, df[extra_cols].iterrows()):
                d = {}
                for col in extra_cols:
                    d[col] = values_serie[col]
                metadatas.append(
                    RawCollarMetadata(hole_id=collar.hole_id, extra_cols=d)
                )

            self._openlog_connection.get_write_iface().import_collar(collars)

            self._openlog_connection.get_write_iface().import_collar_metadata(metadatas)

    def validatePage(self) -> bool:
        """
        Validate current page content (return always True since data is optional)

        Returns: True

        """
        valid = self.dataset_edit.data_is_valid
        df = self.dataset_edit.get_dataframe()
        if df is not None and not self.dataset_edit.crs().isValid():
            valid = False
            QMessageBox.warning(
                self, self.tr("No CRS defined"), self.tr("Define imported data CRS.")
            )

        self.dataset_edit.save_setting(BASE_SETTINGS_KEY)

        return valid
