from geoalchemy2 import WKTElement
from qgis.core import QgsCoordinateReferenceSystem, 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], import_policy: dict = None
    ) -> 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()]

        # store newly created columns. Since rows are inserted just after collar insertion, we need this information to apply import policy.
        created_col_names = []
        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()
                created_col_names.append(col_name)

        # 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():
                is_new_col_name = col_name in created_col_names
                # dont apply policy
                if is_new_col_name:
                    setattr(obj, col_name, value)
                else:
                    policy = import_policy.get(col_name)
                    if policy == "Overwrite":
                        setattr(obj, col_name, value)
                    elif policy == "Merge from DB":
                        if getattr(obj, col_name) is None and value is not None:
                            setattr(obj, col_name, value)
                    elif policy == "Merge to DB":
                        if value is not None:
                            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,
            planned_x=collar.planned_x,
            planned_y=collar.planned_y,
            planned_z=collar.planned_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))
        planned_loc = QgsGeometry(
            QgsPoint(collar.planned_x, collar.planned_y, collar.planned_z)
        )
        # for proj_* geometries, if srid is spherical then 3857 else srid
        is_geographic = QgsCoordinateReferenceSystem(collar.srid).isGeographic()
        if not is_geographic:
            collar.proj_geom = WKTElement(geom.asWkt().replace("Z", ""), srid=-1)
            collar.proj_planned_loc = WKTElement(
                planned_loc.asWkt().replace("Z", ""), srid=-1
            )
        if collar.srid != DEFAULT_SRID:
            geom = convert_to_db_crs(collar.srid, geom)
            planned_loc = convert_to_db_crs(collar.srid, planned_loc)
        # Need to remove Z because of geoalchemy
        collar.geom = WKTElement(geom.asWkt().replace("Z", ""), srid=DEFAULT_SRID)
        collar.planned_loc = WKTElement(
            planned_loc.asWkt().replace("Z", ""), srid=DEFAULT_SRID
        )
        if is_geographic:
            collar.proj_geom = WKTElement(geom.asWkt().replace("Z", ""), srid=-1)
            collar.proj_planned_loc = WKTElement(
                planned_loc.asWkt().replace("Z", ""), srid=-1
            )

        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()

    def selected_collar_edition_available(self):

        return True

    def update_collars(self, collar_list, import_policy: dict = None):
        """
        Override for manually updating geometries.
        """
        # update table
        super().update_collars(collar_list, import_policy)

        # update point geometries
        for collar in collar_list:
            c = (
                self.session.query(self.collar_base)
                .filter(self.collar_base.hole_id == collar.hole_id)
                .first()
            )
            geom = QgsGeometry(QgsPoint(c.x, c.y, c.z))
            planned_loc = QgsGeometry(QgsPoint(c.planned_x, c.planned_y, c.planned_z))

            is_geographic = QgsCoordinateReferenceSystem(c.srid).isGeographic()
            if not is_geographic:
                c.proj_geom = WKTElement(geom.asWkt().replace("Z", ""), srid=-1)
                c.proj_planned_loc = WKTElement(
                    planned_loc.asWkt().replace("Z", ""), srid=-1
                )

            if c.srid != DEFAULT_SRID:
                geom = convert_to_db_crs(c.srid, geom)
                planned_loc = convert_to_db_crs(c.srid, planned_loc)
            c.geom = WKTElement(geom.asWkt().replace("Z", ""), srid=DEFAULT_SRID)
            c.planned_loc = WKTElement(
                planned_loc.asWkt().replace("Z", ""), srid=DEFAULT_SRID
            )

            if is_geographic:
                c.proj_geom = WKTElement(geom.asWkt().replace("Z", ""), srid=-1)
                c.proj_planned_loc = WKTElement(
                    planned_loc.asWkt().replace("Z", ""), srid=-1
                )

        self.session.commit()

    def remove_collars(self, collar_ids):
        """
        Remove collar from database by its ID.
        """
        try:
            for collar_id in collar_ids:
                # remove surveys
                surveys = self._read_iface.get_surveys_from_collars([collar_id])
                for survey in surveys:
                    self.session.delete(survey)
                # remove collar
                self.session.query(self.collar_base).filter(
                    self.collar_base.hole_id == collar_id
                ).delete()

            self.session.commit()
            return True
        except Exception as exc:
            self.session.rollback()
            return False
