from geoalchemy2 import WKTElement
from qgis.core import QgsGeometry, QgsPoint
from sqlalchemy.orm import Session, declarative_base
from xplordb.datamodel.collar import Collar
from xplordb.datamodel.metadata import RawCollarMetadata
from xplordb.datamodel.survey import Survey

from openlog.core.trace_splitter import SpatialiteSplitTracesQueries
from openlog.datamodel.connection.interfaces.categories_interface import (
    CategoriesInterface,
)
from openlog.datamodel.connection.interfaces.read_interface import ReadInterface
from openlog.datamodel.connection.spatialite.database_object import (
    DEFAULT_SRID,
    convert_to_db_crs,
    create_spatialite_collar_metadata_table,
)
from openlog.datamodel.connection.sqlalchemy.sqlalchemy_write_interface import (
    SqlAlchemyWriteInterface,
)

Base = declarative_base()


class SpatialiteWriteInterface(SqlAlchemyWriteInterface):
    def __init__(
        self,
        session: Session,
        person_base: Base,
        dataset_base: Base,
        collar_base: Base,
        survey_base: Base,
        metadata_base: Base,
        lith_base: Base = None,
        read_iface: ReadInterface = None,
        categories_iface: CategoriesInterface = None,
    ):
        """
        Implement WriteInterface with a spatialite session and description of table

        Args:
            session: sqlalchemy session created from engine
            person_base: sqlalchemy base describing person object
            dataset_base: sqlalchemy base describing dataset object
            collar_base: sqlalchemy base describing collar object
            survey_base: sqlalchemy base describing survey object
            lith_base: sqlalchemy base describing lith object
        """
        super().__init__(
            session=session,
            person_base=person_base,
            dataset_base=dataset_base,
            collar_base=collar_base,
            survey_base=survey_base,
            metadata_base=metadata_base,
            lith_base=lith_base,
            categories_iface=categories_iface,
        )
        self._read_iface = read_iface
        self.trace_splitter = SpatialiteSplitTracesQueries

    def import_collar_metadata(self, metadatas: [RawCollarMetadata]) -> None:
        """
        Import metadata into OpenLogConnection interface.
        Table is altered first, then values are updated.

        Args:
            metadatas: collar metadatas to be imported
        """

        # add column if dont exists in database

        extra_cols = {}
        for m in metadatas:
            for col_name, value in m.extra_cols.items():
                sql_type = "FLOAT" if type(value) != str else "TEXT"
                extra_cols[col_name] = sql_type

        q = """
            select name from pragma_table_info('metadata') as tblInfo;
            """
        col_base = [r[0] for r in self.session.execute(q).fetchall()]

        for col_name, col_type in extra_cols.items():

            if col_name not in col_base:
                q = f"ALTER TABLE metadata ADD COLUMN {col_name} {col_type} ;"
                self.session.execute(q)
                self.session.commit()

        # update data
        for metadata in metadatas:
            # generate an updated SpatialiteCollar definition with extra columns
            _, cls = create_spatialite_collar_metadata_table(m)
            obj = self.session.query(cls).get(metadata.hole_id)
            for col_name, value in metadata.extra_cols.items():
                setattr(obj, col_name, value)

        self.session.commit()

    def _create_collar_base(self, collar: Collar):
        """
        Override collar sqlalchemy Base object creation to add geometry column definition from x / y / z

        Args:
            collar: Collar object

        Returns: sqlalchemy base describing collar object

        """
        collar = self.collar_base(
            hole_id=collar.hole_id,
            data_set=collar.data_set,
            loaded_by=collar.loaded_by,
            x=collar.x,
            y=collar.y,
            z=collar.z,
            srid=collar.srid,
            project_srid=collar.project_srid,
            eoh=collar.eoh,
            dip=collar.dip,
            azimuth=collar.azimuth,
            planned_eoh=collar.planned_eoh,
            survey_date=collar.survey_date,
        )
        geom = QgsGeometry(QgsPoint(collar.x, collar.y, collar.z))
        if collar.srid != DEFAULT_SRID:
            geom = convert_to_db_crs(collar.srid, geom)
        # Need to remove Z because of geoalchemy
        collar.geom = WKTElement(geom.asWkt().replace("Z", ""), srid=DEFAULT_SRID)
        return collar

    def selected_collar_survey_definition_available(self) -> bool:
        """
        Return True if the collar survey definition is available
        :return: A boolean value.
        """
        return True

    def replace_collar_surveys(self, collar_id: str, surveys: [Survey]) -> None:
        """
        Replace survey for collar

        Args:
            collar_id:: The collar id of the hole
            surveys: new survey for collar
        """
        current_survey = self._read_iface.get_surveys_from_collars([collar_id])
        for survey in current_survey:
            self.session.delete(survey)
        self.session.commit()
        self.import_surveys(surveys)

    def replace_collar_planned_surveys(self, collar_id: str, survey: Survey) -> None:

        """
        Replace planned survey for collar

        Args:
            collar_id:: The collar id of the hole
            survey: new planned survey for collar
        """
        collar = self._read_iface.get_collar(collar_id)
        collar.dip = survey.dip
        collar.azimuth = survey.azimuth
        collar.planned_eoh = survey.depth
        self.session.commit()
