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], import_policy: dict = None) -> None:
        """
        Import collar into OpenLogConnection interface.

        Args:
            collars: collars to be imported
        """
        existing_collars = self.session.query(self.collar_base).all()
        existing_collars_id = [c.hole_id for c in existing_collars]
        to_insert = [
            collar for collar in collars if collar.hole_id not in existing_collars_id
        ]
        to_update = [
            collar for collar in collars if collar.hole_id in existing_collars_id
        ]

        for collar in to_insert:
            new_collar = self._create_collar_base(collar)
            self.session.add(new_collar)

        if len(to_update) > 0:
            self.update_collars(to_update, import_policy)

    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], import_policy: dict = None) -> None:
        """
        Import survey into OpenLogConnection interface.

        Raises OpenLogConnection.ImportException on import failure

        Args:
            surveys: surveys to be imported
        """
        existing_surveys = self.session.query(self.survey_base).all()
        existing_surveys_id = [(c.hole_id, c.depth) for c in existing_surveys]
        to_insert = [
            survey
            for survey in surveys
            if (survey.hole_id, survey.depth) not in existing_surveys_id
        ]
        to_update = [
            survey
            for survey in surveys
            if (survey.hole_id, survey.depth) in existing_surveys_id
        ]

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

        self.update_surveys(to_update, import_policy)

    def update_surveys(self, surveys: [Survey], import_policy: dict = None) -> None:
        """
        Update existing surveys according to import policy.

        Raises OpenLogConnection.ImportException on import failure

        Args:
            surveys: surveys to be updated
        """
        updated_surveys = surveys
        # update values according to update policies
        if import_policy is not None:
            updated_surveys = []
            # surveys come from import wizard
            for survey in surveys:
                existing_survey = (
                    self.session.query(self.survey_base)
                    .filter(
                        self.survey_base.hole_id == survey.hole_id,
                        self.survey_base.depth == survey.depth,
                    )
                    .first()
                )
                for col, policy in import_policy.items():
                    if policy == "Overwrite":
                        setattr(existing_survey, col, getattr(survey, col))
                    elif policy == "Merge from DB":
                        if getattr(existing_survey, col) is None:
                            setattr(existing_survey, col, getattr(survey, col))
                    elif policy == "Merge to DB":
                        if getattr(survey, col) is not None:
                            setattr(existing_survey, col, getattr(survey, col))

                updated_surveys.append(existing_survey)

        try:
            for survey in updated_surveys:
                self.session.query(self.survey_base).filter(
                    self.survey_base.hole_id == survey.hole_id,
                    self.survey_base.depth == survey.depth,
                ).update({"dip": survey.dip, "azimuth": survey.azimuth})

            self.session.commit()

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

    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, import_policy: dict = None) -> None:
        """
        Update collars from database.
        Args:
            - collar_list : list of collars to update
            - import_policy : dictionnary
        """
        updated_collars = collar_list
        # update values according to update policies
        if import_policy is not None:
            updated_collars = []
            # collars come from import wizard
            for collar in collar_list:
                existing_collar = (
                    self.session.query(self.collar_base)
                    .filter(self.collar_base.hole_id == collar.hole_id)
                    .first()
                )
                for col, policy in import_policy.items():
                    # exclude metadata
                    if not hasattr(existing_collar, col):
                        continue
                    if policy == "Overwrite":
                        setattr(existing_collar, col, getattr(collar, col))
                    elif policy == "Merge from DB":
                        if getattr(existing_collar, col) is None:
                            setattr(existing_collar, col, getattr(collar, col))
                    elif policy == "Merge to DB":
                        if getattr(collar, col) is not None:
                            setattr(existing_collar, col, getattr(collar, col))

                updated_collars.append(existing_collar)

        try:
            for collar in updated_collars:
                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)
