from sqlalchemy.orm import Session, declarative_base
from xplordb.datamodel.collar import Collar
from xplordb.datamodel.lith import Lith
from xplordb.datamodel.survey import Survey

from openlog.datamodel.connection.interfaces.read_interface import ReadInterface

Base = declarative_base()


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

        self.person_base = person_base
        self.dataset_base = dataset_base
        self.collar_base = collar_base
        self.survey_base = survey_base
        self.lith_base = lith_base

    def get_available_person_codes(self) -> [str]:
        """
        Return available person codes from connection

        Returns: available person codes in xplordb connection

        """
        persons = self.session.query(self.person_base).all()
        return [person.code for person in persons]

    def get_available_dataset_names(self) -> [str]:
        """
        Return available dataset names from connection

        Returns: available dataset names in xplordb connection

        """

        datasets = self.session.query(self.dataset_base).all()
        return [dataset.name for dataset in datasets]

    def get_collar(self, hole_id: str) -> Collar:
        """
        The function `get_collar` returns a collar object for a given hole_id

        :param hole_id: The hole_id of the collar to be updated
        :type hole_id: str
        :return: A Collar object
        """
        return (
            self.session.query(self.collar_base)
            .filter(self.collar_base.hole_id == hole_id)
            .first()
        )

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

        """
        if planned:
            collars = [self.get_collar(id) for id in collars_id]
            surveys = [
                Survey(
                    hole_id=c.hole_id,
                    data_set=c.data_set,
                    loaded_by=c.loaded_by,
                    depth=0.0,
                    dip=c.dip,
                    azimuth=c.azimuth,
                )
                for c in collars
            ]
        else:
            surveys = (
                self.session.query(self.survey_base)
                .filter(self.survey_base.hole_id.in_(collars_id))
                .all()
            )

        return surveys
