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]) -> [Survey]:
        """
        Return surveys from collar id list

        """
        return (
            self.session.query(self.survey_base)
            .filter(self.survey_base.hole_id.in_(collars_id))
            .all()
        )
