from sqlalchemy.orm import Session, declarative_base
from xplordb.datamodel.collar import Collar
from xplordb.datamodel.dataset import Dataset
from xplordb.datamodel.lith import Lith
from xplordb.datamodel.metadata import RawCollarMetadata
from xplordb.datamodel.person import Person
from xplordb.datamodel.survey import Survey

from openlog.core.trace_splitter import SplitTracesQueries
from openlog.datamodel.assay.generic_assay import (
    AssayDatabaseDefinition,
    AssayDataExtent,
)
from openlog.datamodel.connection.interfaces.categories_interface import (
    CategoriesInterface,
)
from openlog.datamodel.connection.interfaces.write_interface import WriteInterface

Base = declarative_base()


class SqlAlchemyWriteInterface(WriteInterface):
    def __init__(
        self,
        session: Session,
        person_base: Base,
        dataset_base: Base,
        collar_base: Base,
        survey_base: Base,
        metadata_base: Base,
        lith_base: Base = None,
        categories_iface: CategoriesInterface = None,
    ):
        """
        Implement WriteInterface 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
            categories_iface: CategoriesInterface for category insertion for liths
        """
        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
        self.metadata_base = metadata_base
        self._categories_iface = categories_iface
        self.trace_splitter = SplitTracesQueries

    def can_import_collar(self) -> bool:
        """
        Indicate if connection can import collar

        Returns: True

        """
        return True

    def import_collar(self, collars: [Collar]) -> None:
        """
        Import collar into OpenLogConnection interface.

        Args:
            collars: collars to be imported
        """
        for collar in collars:
            new_collar = self._create_collar_base(collar)
            self.session.add(new_collar)

    def import_datasets(self, datasets: [Dataset]) -> None:
        """
        Import dataset into OpenLogConnection interface.

        Args:
            datasets: datasets to be imported
        """
        for dataset in datasets:
            new_dataset = self._create_dataset_base(dataset)
            self.session.add(new_dataset)

    def import_persons(self, persons: [Person]) -> None:
        """
        Import persons into OpenLogConnection interface.

        Args:
            persons: persons to be imported
        """
        for person in persons:
            new_person = self._create_person_base(person)
            self.session.add(new_person)

    def import_surveys(self, surveys: [Survey]) -> None:
        """
        Import survey into OpenLogConnection interface.

        Raises OpenLogConnection.ImportException on import failure

        Args:
            surveys: surveys to be imported
        """

        for survey in surveys:
            new_survey = self._create_survey_base(survey)
            self.session.add(new_survey)

    def person_code_max_size(self) -> int:
        """
        Define maximum size for person code creation

        Returns: 10

        """
        return 10

    def _create_collar_base(self, collar: Collar) -> Base:
        """
        Create collar Base object from collar

        Args:
            collar: collar object

        Returns: collar base object

        """
        return 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,
        )

    def _create_dataset_base(self, dataset: Dataset) -> Base:
        """
        Create dataset Base object from dataset

        Args:
            dataset: dataset object

        Returns: dataset base object

        """
        return self.dataset_base(
            name=dataset.name, full_name=dataset.full_name, loaded_by=dataset.loaded_by
        )

    def _create_person_base(self, person: Person) -> Base:
        """
        Create person Base object from person

        Args:
            person: person object

        Returns: person base object

        """
        return self.person_base(trigram=person.code, full_name=person.full_name)

    def _create_survey_base(self, survey: Survey):
        """
        Create survey Base object from survey

        Args:
            survey: survey object

        Returns: survey base object

        """
        return self.survey_base(
            hole_id=survey.hole_id,
            data_set=survey.data_set,
            loaded_by=survey.loaded_by,
            depth=survey.depth,
            dip=survey.dip,
            azimuth=survey.azimuth,
        )

    def _create_lith_base(self, lith: Lith) -> Base:
        """
        Create lith Base object from lith

        Args:
            lith: Lith object

        Returns: lith base object

        """
        if self.lith_base:
            return self.lith_base(
                hole_id=lith.hole_id,
                data_set=lith.data_set,
                lith_code=lith.lith_code,
                from_m=lith.from_m,
                to_m=lith.to_m,
                logged_by=lith.logged_by,
                loaded_by=lith.loaded_by,
            )
        else:
            return None

    def update_collars(self, collar_list) -> None:
        """
        Update collars from database.
        Args:
            - collar_list : list of collars to update
        """
        try:
            for collar in collar_list:
                self.session.query(self.collar_base).filter(
                    self.collar_base.hole_id == collar.hole_id
                ).update(
                    {
                        "x": collar.x,
                        "y": collar.y,
                        "z": collar.z,
                        "planned_x": collar.planned_x,
                        "planned_y": collar.planned_y,
                        "planned_z": collar.planned_z,
                        "eoh": collar.eoh,
                        "planned_eoh": collar.planned_eoh,
                        "srid": collar.srid,
                    }
                )

            self.session.commit()

        except Exception as exc:
            self.session.rollback()
            raise WriteInterface.ImportException(exc)
