import pandas as pd
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 xplordb.import_data import ImportData

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


class XplordbWriteInterface(WriteInterface):
    def __init__(self, session, read_iface: ReadInterface):
        """
        Implement WriteInterface with a xplordb session

        Args:
            session: sqlalchemy session created from engine
        """
        super(XplordbWriteInterface, self).__init__()
        self.session = session
        self._read_iface = read_iface

    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 xplordb database.

        Raises WriteInterface.ImportException on import failure

        Args:
            collars: collars to be imported
        """
        existing_collars = self._read_iface.get_all_collars()
        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
        ]
        try:
            import_data = ImportData(self.session)
            import_data.import_collars_array(to_insert)
            if len(to_update) > 0:
                self.update_collars(to_update, import_policy)

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

    def _update_collar_metadata(
        self, metadatas: [RawCollarMetadata], import_policy: dict = None
    ) -> list[RawCollarMetadata]:
        """
        Update RawCollarMetadata according to import policy.
        """
        # get existing metadata table
        if len(metadatas) == 0:
            return []

        # get current table columns
        q = "select column_name from information_schema.columns WHERE table_name = 'metadata' AND table_schema = 'dh'"
        table_columns = self.session.execute(q).fetchall()
        table_columns = [col[0] for col in table_columns]

        # load existing data as dict
        data = self.session.execute("select * from dh.metadata").fetchall()
        df = pd.DataFrame.from_records(data, columns=table_columns)
        df = df.set_index("hole_id")
        actual = df.to_dict("index")

        # update each RawCollarMetadata
        for metadata in metadatas:
            for col, new_value in metadata.extra_cols.items():
                if col not in table_columns:
                    continue

                policy = import_policy.get(col)
                # check hole id is present
                if actual.get(metadata.hole_id) is None:
                    continue
                db_value = actual.get(metadata.hole_id).get(col)
                if policy == "Overwrite":
                    metadata.extra_cols[col] = new_value

                elif policy == "Merge from DB":
                    if db_value is None and new_value is not None:
                        metadata.extra_cols[col] = new_value

                elif policy == "Merge to DB":
                    if new_value is not None:
                        metadata.extra_cols[col] = new_value

                elif policy == "Ignore":
                    metadata.extra_cols[col] = db_value

        return metadatas

    def import_collar_metadata(
        self, metadatas: [RawCollarMetadata], import_policy: dict = None
    ) -> None:
        """
        Import collar into xplordb database.

        Raises WriteInterface.ImportException on import failure

        Args:
            collars: collars to be imported
        """
        try:
            updated_metadatas = self._update_collar_metadata(metadatas, import_policy)
            import_data = ImportData(self.session)
            import_data.import_metadatas_array(updated_metadatas)
        except ImportData.ImportException as exc:
            self.session.rollback()
            raise WriteInterface.ImportException(exc)

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

        Raises WriteInterface.ImportException on import failure

        Args:
            datasets: datasets to be imported
        """
        try:
            import_data = ImportData(self.session)
            import_data.import_datasets_array(datasets)
        except ImportData.ImportException as exc:
            self.session.rollback()
            raise WriteInterface.ImportException(exc)

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

        Returns: 3

        """
        return 3

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

        Raises WriteInterface.ImportException on import failure

        Args:
            persons: persons to be imported
        """
        try:
            import_data = ImportData(self.session)
            import_data.import_persons_array(persons)
        except ImportData.ImportException as exc:
            self.session.rollback()
            raise WriteInterface.ImportException(exc)

    def import_surveys(self, surveys: [Survey], import_policy: dict = None) -> None:
        """
        Import survey into xplordb database.

        Raises WriteInterface.ImportException on import failure

        Args:
            surveys: surveys to be imported
        """
        existing_surveys = self._read_iface.get_all_surveys()
        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
        ]
        try:
            import_data = ImportData(self.session)
            import_data.import_surveys_array(to_insert)
            if len(to_update) > 0:
                self.update_surveys(to_update, import_policy)

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

    def selected_collar_survey_definition_available(self) -> bool:
        """
        Return True if the collar survey definition is available
        :return: A boolean value.
        """
        return True

    def selected_collar_edition_available(self):

        return True

    def replace_collar_surveys(self, collar_id: str, surveys: [Survey]) -> None:
        """
        Replace survey for collar

        Args:
            collar_id:: The collar id of the hole
            surveys: new survey for collar
        """
        current_survey = self._read_iface.get_surveys_from_collars([collar_id])
        for survey in current_survey:
            self.session.delete(survey)
        self.session.flush()
        self.import_surveys(surveys)

    def replace_collar_planned_surveys(self, collar_id: str, survey: Survey) -> None:

        """
        Replace planned survey for collar

        Args:
            collar_id:: The collar id of the hole
            survey: new planned survey for collar
        """
        collar = self._read_iface.get_collar(collar_id)
        collar.dip = survey.dip
        collar.azimuth = survey.azimuth
        collar.planned_eoh = survey.depth
        self.session.commit()

    def update_surveys(self, survey_list, import_policy: dict = None) -> None:
        """
        Update surveys from database.
        Args:
            - survey_list : list of surveys to update
        """

        try:

            # update values according to update policies
            if import_policy is not None:
                updated_surveys = []
                # surveys come from import wizard
                for survey in survey_list:
                    existing_survey = self._read_iface.get_survey(
                        hole_id=survey.hole_id, depth=survey.depth
                    )
                    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)

                self.session.commit()
            else:

                for survey in survey_list:
                    db_survey = self._read_iface.get_survey(
                        hole_id=survey.hole_id, depth=survey.depth
                    )
                    db_survey.dip = (survey.dip,)
                    db_survey.azimuth = (survey.azimuth,)

                self.session.commit()

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

    def update_collars(self, collar_list, import_policy: dict = None) -> None:
        """
        Update collars from database.
        Args:
            - collar_list : list of collars to update
        """

        try:

            # 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._read_iface.get_collar(
                        hole_id=collar.hole_id
                    )
                    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)

                self.session.commit()
            else:

                for collar in collar_list:
                    db_collar = self._read_iface.get_collar(hole_id=collar.hole_id)
                    db_collar.x = (collar.x,)
                    db_collar.y = collar.y
                    db_collar.z = (collar.z,)
                    db_collar.planned_x = (collar.planned_x,)
                    db_collar.planned_y = (collar.planned_y,)
                    db_collar.planned_z = (collar.planned_z,)
                    db_collar.eoh = (collar.eoh,)
                    db_collar.planned_eoh = (collar.planned_eoh,)
                    db_collar.srid = (collar.srid,)

                self.session.commit()

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

    def remove_collars(self, collar_ids):
        """
        Remove collar from database by its ID.
        """
        try:
            for collar_id in collar_ids:
                # first, delete surveys
                self.session.execute(
                    f"DELETE FROM dh.surv WHERE hole_id = '{collar_id}';"
                )

                # delete metadatas
                self.session.execute(
                    f"DELETE FROM dh.metadata WHERE hole_id = '{collar_id}';"
                )

                collar = self._read_iface.get_collar(collar_id)
                self.session.delete(collar)

            self.session.commit()
            return True
        except Exception as exc:
            self.session.rollback()
            return False
