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.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,
)
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,
        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,
            lith_base=lith_base,
            categories_iface=categories_iface,
        )
        self._read_iface = read_iface
        self.trace_splitter = SpatialiteSplitTracesQueries

    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,
            eoh=collar.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)
