from sqlalchemy.orm import Session, declarative_base
from xplordb.datamodel.survey import Survey

from openlog.datamodel.connection.sqlalchemy.sqlalchemy_read_interface import (
    SqlAlchemyReadInterface,
)

Base = declarative_base()


class GeoticReadInterface(SqlAlchemyReadInterface):
    def __init__(
        self,
        session: Session,
        person_base: Base,
        dataset_base: Base,
        collar_base: Base,
        survey_base: Base,
        lith_base: Base = None,
    ):
        """
        Implement ReadInterface with an Geotic 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,
        )

    def get_surveys_from_collars(self, collars_id: [str]) -> [Survey]:
        """
        Return surveys from collar id list

        """
        # Conversion from hole_id string to Geotic id_collar
        collars = (
            self.session.query(self.collar_base)
            .filter(self.collar_base.hole_id.in_(collars_id))
            .all()
        )
        id_key = [collar_base.idKey for collar_base in collars]

        # Only use survey with valid azimuth and dip
        surveys = (
            self.session.query(self.survey_base)
            .filter(
                self.survey_base.id_collar.in_(id_key),
                self.survey_base.invalid_azimuth == 0,
                self.survey_base.invalid_dip == 0,
                self.survey_base.azimuth.is_not(None),
                self.survey_base.dip.is_not(None),
            )
            .all()
        )
        # Add first survey
        for collar_base in collars:
            # Only collar with valid coordinates
            if (
                collar_base.coord
                and collar_base.coord[0].azimuth is not None
                and collar_base.coord[0].dip is not None
            ):
                surveys.append(
                    Survey(
                        hole_id=collar_base.hole_id,
                        data_set=collar_base.data_set,
                        loaded_by=collar_base.loaded_by,
                        depth=0.0,
                        azimuth=collar_base.coord[0].azimuth,
                        dip=collar_base.coord[0].dip,
                    )
                )
        return surveys
