import json
from typing import Any, List

import numpy as np
import pandas as pd
import sqlalchemy
from sqlalchemy.orm import declarative_base

from openlog.core.geo_extractor import GeoExtractor
from openlog.core.trace_splitter import SplitTracesQueries
from openlog.datamodel.assay.generic_assay import (
    AssayColumn,
    AssayDatabaseDefinition,
    AssayDataExtent,
    AssayDefinition,
    AssayDomainType,
    AssaySeriesType,
    GenericAssay,
)
from openlog.datamodel.connection.interfaces.assay_interface import AssayInterface
from openlog.datamodel.connection.openlog_connection import OpenLogConnection
from openlog.datamodel.connection.sqlalchemy.assay_factory import AssayFactory
from openlog.datamodel.connection.sqlalchemy.assay_tables import (
    create_assay_column_table,
    create_assay_database_column_definition_base,
    create_assay_database_definition_base,
    create_assay_definition_base,
)
from openlog.toolbelt import PlgLogger


class SqlAlchemyAssayInterface(AssayInterface):
    """
    sqlalchemy session use for OpenLogConnection interface implementation.

    Here are the implemented interfaces :

    - create assay in connection : add_assay_table
    - get available assay and assay database definition : get_all_available_assay_definitions / get_assay_table_definition
    - insert assay into connection : import_assay_data
    - get assay from connection : get_assay
    - check for feature availability : can_import_assay / can_import_assay_data

    This interface is used by :
    - SpatialiteConnection
    - XplordbConnection

    Some object describing database must be defined in constructor:

    self.session : sqlalchemy session created with wanted engine (spatialite, mssql, ...)
    self.assay_definition_base : sqlalchemy Base object describing assay definition access
    self.assay_database_definition_base : sqlalchemy Base object describing assay database definition access

    """

    def __init__(self, engine, session, schema: str = ""):
        super().__init__()

        self.log = PlgLogger().log
        self.engine = engine
        self.session = session
        self.schema = schema
        self.trace_splitter = SplitTracesQueries
        self.geo_extractor = GeoExtractor

        AssayBase = declarative_base()

        self.assay_definition_base = create_assay_definition_base(AssayBase, schema)
        self.assay_column_definition_base = create_assay_column_table(AssayBase, schema)

        self.assay_database_definition_base = create_assay_database_definition_base(
            AssayBase, schema
        )
        self.assay_column_database_definition_base = (
            create_assay_database_column_definition_base(AssayBase, schema)
        )

        self._create_tables()

        # Maps to get assay base and assay database definition object
        # Define sqlalchemy base for available assays
        try:
            # we must give it a chance to be able to update database schema
            self.assay_database_definition_map = self._init_assay_maps()
        except (sqlalchemy.exc.OperationalError, sqlalchemy.exc.ProgrammingError):
            pass

        # List of table added for assay : need to be dropped from database in case of rollback
        self._session_added_tables = []

    def rollback(self) -> None:

        """
        Rollback current changes

        Raises OpenLogConnection.ImportException if rollback fail

        """
        try:
            for table in self._session_added_tables:
                table.drop(self.engine)
            self._session_added_tables.clear()

        except sqlalchemy.exc.SQLAlchemyError as exc:
            raise OpenLogConnection.ImportException(exc)

    def commit(self) -> None:

        """
        Commit current changes

        Raises OpenLogConnection.ImportException on import failure

        """
        self._session_added_tables.clear()

    def _create_tables(self) -> None:
        try:
            self._before_sqlalchemy_table_creation(
                self.assay_definition_base.__table__, self.schema
            )
            self.assay_definition_base.__table__.create(self.engine)
            self._before_sqlalchemy_table_creation(
                self.assay_column_definition_base.__table__, self.schema
            )
            self.assay_column_definition_base.__table__.create(self.engine)

            self._before_sqlalchemy_table_creation(
                self.assay_database_definition_base.__table__, self.schema
            )
            self.assay_database_definition_base.__table__.create(self.engine)
            self._before_sqlalchemy_table_creation(
                self.assay_column_database_definition_base.__table__, self.schema
            )
            self.assay_column_database_definition_base.__table__.create(self.engine)

        except (
            sqlalchemy.exc.ProgrammingError,
            sqlalchemy.exc.OperationalError,
        ) as exc:
            # No error if tables already exists
            self.log(
                message="Exception raised for assay tables creation. It must come from already created "
            )

    def _init_assay_maps(self) -> {str: AssayDatabaseDefinition}:
        """
        Read all available assay definitions and create a map of sqlalchemy base

        Returns: map of sqlalchemy base

        """
        assay_database_definition_map = {}

        assay_definition_base = self.get_all_available_assay_definitions()
        for assay_definition in assay_definition_base:

            assay_table = (
                self.session.query(self.assay_database_definition_base)
                .filter(
                    self.assay_database_definition_base.variable
                    == assay_definition.variable
                )
                .first()
            )
            assay_database_definition = assay_table.to_assay_database_definition()

            assay_database_definition_map[
                assay_definition.variable
            ] = assay_database_definition

        return assay_database_definition_map

    def can_import_assay(self) -> bool:
        """
        Return True if connection can add assay table
        :return: True
        """
        return True

    def add_assay_table(
        self, assay_definition: AssayDefinition, assay_table: AssayDatabaseDefinition
    ) -> None:
        """
        Add a table in database for assay import

        Args:
            assay_definition: AssayDefinition used to describe assay contents
            assay_table: AssayDatabaseDefinition  used to decribe how to insert assay into database
        """
        # Add assay definition in assays table
        try:
            new_table = self.assay_definition_base()
            new_table.from_assay_definition(assay_definition)
            self.session.add(new_table)

            new_table_definition = self.assay_database_definition_base()
            new_table_definition.from_assay_database_definition(
                assay_definition.variable, assay_table
            )
            new_table_definition.variable = assay_definition.variable
            self.session.add(new_table_definition)

        except sqlalchemy.exc.IntegrityError as exc:
            self.session.rollback()
            raise AssayInterface.ImportException(exc)

        # Create assay table
        try:
            self.assay_database_definition_map[assay_definition.variable] = assay_table

            sql_alchemy_table = AssayFactory(
                assay_definition, assay_table
            ).create_assay_table()

            sql_alchemy_table = self._cast_column_type(
                sql_alchemy_table, assay_definition
            )

            self._before_sqlalchemy_table_creation(
                sql_alchemy_table, assay_table.schema
            )
            sql_alchemy_table.create(self.engine)
            self._session_added_tables.append(sql_alchemy_table)
        except (
            sqlalchemy.exc.InvalidRequestError,
            sqlalchemy.exc.OperationalError,
            sqlalchemy.exc.ProgrammingError,
        ) as exc:
            self.session.rollback()
            raise AssayInterface.ImportException(exc)

    def get_all_available_assay_definitions(self) -> List[AssayDefinition]:
        """
        Return all assay definitions available in connection

        Return empty list by default must be implemented in interface

        """
        assay_def = []
        try:
            assay_tables = self.session.query(self.assay_definition_base).all()
            assay_def = [
                assay_table.to_assay_definition() for assay_table in assay_tables
            ]
        except:
            self.session.rollback()
            pass

        return assay_def

    def get_assay_table_definition(
        self, assay_variable: str
    ) -> AssayDatabaseDefinition:
        """
        Get AssayDatabaseDefinition from assay variable

        Args:
            assay_variable: assay variable name

        Returns:
            AssayDatabaseDefinition: assay database definition
        """
        self.assay_database_definition_map = self._init_assay_maps()
        return self.assay_database_definition_map[assay_variable]

    def can_import_assay_data(self) -> bool:
        """
        Return True if connection can import assay data
        :return: True

        """
        return True

    def import_assay_data(
        self, assays: List[GenericAssay], import_policy: dict = {}
    ) -> None:
        """
        Import assay into connection

        raises OpenLogConnection.InvalidInterface if not implemented

        Args:
            assays: assay list to be imported
        """
        for assay in assays:
            self._check_assay(assay)
            self._import_assay(assay, import_policy)

        # create a layer for each column
        if len(assays) == 0:
            return

        assay = assays[0]
        if assay.assay_definition.domain == AssayDomainType.DEPTH:

            for assay_column in assay.assay_definition.columns.values():
                if assay_column.series_type not in (
                    AssaySeriesType.NUMERICAL,
                    AssaySeriesType.CATEGORICAL,
                ):
                    continue
                # Traces
                self.split_traces(
                    table_name=assay.assay_definition.variable,
                    assay_column=assay_column,
                    discrete=True
                    if assay.assay_definition.data_extent == AssayDataExtent.DISCRETE
                    else False,
                    geom="proj_effective_geom",
                )
                # planned traces
                self.split_traces(
                    table_name=assay.assay_definition.variable,
                    assay_column=assay_column,
                    discrete=True
                    if assay.assay_definition.data_extent == AssayDataExtent.DISCRETE
                    else False,
                    geom="proj_planned_geom",
                )

    @staticmethod
    def _check_assay(assay: GenericAssay) -> None:
        """
        Check assay data
        Raise AssayInterface.ImportException in case of invalid data

        Args:
            assay: (GenericAssay) assay to be imported
        """
        try:
            assay.check_import_data()
        except GenericAssay.InvalidImportData as exc:
            raise AssayInterface.ImportException(exc)

    def _import_assay(self, assay: GenericAssay, import_policy: dict = {}) -> None:
        """
        Import assay into OpenLogConnection

        Args:
            assay: (GenericAssay) assay to be imported
            import_policy : actions to do if data already exists in database.
        """
        # Get base sqlalchemy class for data import
        assay_table_definition = self.get_assay_table_definition(
            assay.assay_definition.variable
        )
        assay_table = AssayFactory(
            assay.assay_definition, assay_table_definition
        ).create_assay_table()

        assay_table = self._cast_column_type(assay_table, assay.assay_definition)

        params = []
        for i in range(0, len(assay.import_x_values)):
            value = {
                assay_table_definition.hole_id_col: assay.hole_id,
                assay_table_definition.dataset_col: assay.dataset,
                assay_table_definition.person_col: assay.person,
                assay_table_definition.import_date_col: assay.import_date,
            }

            if assay.assay_definition.data_extent == AssayDataExtent.DISCRETE:
                value[assay_table_definition.x_col] = self._convert_if_numpy_value(
                    assay.import_x_values[i]
                )
            else:
                value[assay_table_definition.x_col] = self._convert_if_numpy_value(
                    assay.import_x_values[i][0]
                )
                value[assay_table_definition.x_end_col] = self._convert_if_numpy_value(
                    assay.import_x_values[i][1]
                )

            for col, vals in assay.import_y_values.items():
                col_used = (
                    assay_table_definition.y_col[col]
                    if col in assay_table_definition.y_col
                    else col
                )
                value[col_used] = self._convert_if_numpy_value(vals[i])

            params.append(value)

        params = self._convert_composite_string(params, assay.assay_definition)

        try:
            # check if assay already exists
            assay_db = self.get_assay(
                hole_id=assay.hole_id, variable=assay.assay_definition.variable
            )
            assay_db_values = {}
            for col in assay_db.assay_definition.columns.keys():
                x, y = assay_db.get_all_values(col)
                assay_db_values["x"] = x
                assay_db_values[col] = y
                # add uncertainties or detection limits that needs to be updated
                assay_column = assay_db.assay_definition.columns[col]
                for uncert_col in assay_column.uncertainty.get_uncertainty_columns():
                    _, y = assay_db.get_all_values(uncert_col)
                    assay_db_values[uncert_col] = y

                for (
                    detection_col
                ) in assay_column.detection_limit.get_detection_columns():
                    _, y = assay_db.get_all_values(detection_col)
                    assay_db_values[detection_col] = y

            # convert datetime (from db) to TimeStamp
            # print("before conversion",  assay_db_values["x"])
            convert_date_func = np.vectorize(
                lambda x: pd.to_datetime(x).tz_localize(0).tz_convert(None),
                otypes=[object],
            )
            if assay_db.assay_definition.domain == AssayDomainType.TIME:
                assay_db_values["x"] = convert_date_func(assay_db_values["x"])
            else:
                assay_db_values["x"] = assay_db_values["x"].astype(float)

            # primary keys
            if assay_db.assay_definition.data_extent == AssayDataExtent.EXTENDED:
                db_pks = [
                    (
                        assay_db.hole_id,
                        assay_db_values["x"][i, 0],
                        assay_db_values["x"][i, 1],
                    )
                    for i in range(assay_db_values["x"].shape[0])
                ]
                db_params = [
                    {
                        "hole": assay_db.hole_id,
                        "x": assay_db_values["x"][i, 0],
                        "x_end": assay_db_values["x"][i, 1],
                    }
                    | {
                        col: assay_db_values[col][i]
                        for col in assay_db_values.keys()
                        if col != "x"
                    }
                    for i in range(assay_db_values["x"].shape[0])
                ]
                new_pks = [
                    (param["hole"], param["x"], param["x_end"]) for param in params
                ]
            else:
                db_pks = [
                    (assay_db.hole_id, assay_db_values["x"][i])
                    for i in range(assay_db_values["x"].shape[0])
                ]
                new_pks = [(param["hole"], param["x"]) for param in params]
                db_params = [
                    {"hole": assay_db.hole_id, "x": assay_db_values["x"][i]}
                    | {
                        col: assay_db_values[col][i]
                        for col in assay_db_values.keys()
                        if col != "x"
                    }
                    for i in range(assay_db_values["x"].shape[0])
                ]

            # print("db_pks", db_pks, "new_pks", new_pks)
            # get elements to insert (no occurence in database)
            pks_to_insert = [new_pk for new_pk in new_pks if new_pk not in db_pks]
            indexes = [i for i in range(len(new_pks)) if new_pks[i] in pks_to_insert]

            params_to_insert = [params[i] for i in indexes]
            params_to_import = [
                params[i] for i in range(len(params)) if i not in indexes
            ]
            # print("insert", params_to_insert, "import", params_to_import)
            if len(params_to_insert) > 0:
                self.session.execute(assay_table.insert(), params_to_insert)

            # update database data according policies
            ## sort params_to_import and db_params to be paired

            if assay_db.assay_definition.data_extent == AssayDataExtent.EXTENDED:
                sorting_key = lambda x: (x["hole"], x["x"], x["x_end"])
            else:
                sorting_key = lambda x: (x["hole"], x["x"])
            params_to_import = sorted(params_to_import, key=sorting_key)
            db_params = sorted(db_params, key=sorting_key)

            updated_params = []
            column_names = [col for col in assay_db_values.keys() if col != "x"]
            for param_to_import, db_param in zip(params_to_import, db_params):
                # we start from file data as base
                new_param = param_to_import.copy()
                for col in column_names:

                    policy = import_policy.get(col)
                    if policy is None:
                        policy = "Ignore"
                    # replace np.nan by None
                    try:
                        db_param[col] = (
                            None if np.isnan(db_param[col]) else db_param[col]
                        )
                    except TypeError:
                        db_param[col] = (
                            None if db_param[col] == "None" else db_param[col]
                        )

                    # if spherical, convert str data to dict
                    if (
                        assay_db.assay_definition.columns.get(col)
                        and assay_db.assay_definition.columns[col].series_type
                        == AssaySeriesType.SPHERICAL
                        and new_param[col] is not None
                    ):
                        new_param[col] = self._convert_spherical_to_dict(new_param[col])

                    # keep db values
                    if policy == "Ignore":
                        # new_param[col] = db_param[col]
                        new_param.pop(col)
                    # fill database NULL with file values
                    elif policy == "Merge from DB":
                        if db_param[col] is not None:
                            # new_param[col] = db_param[col]
                            new_param.pop(col)
                        elif new_param[col] is None:
                            new_param.pop(col)

                    # fill file NULL with database values
                    elif policy == "Merge to DB":
                        if new_param[col] == db_param[col]:
                            new_param.pop(col)
                            continue

                        if new_param[col] is None:
                            new_param[col] = db_param[col]

                    elif policy == "Overwrite":
                        if new_param[col] == db_param[col]:
                            new_param.pop(col)
                        continue
                # at least one column to update
                if any([col for col in column_names if col in new_param.keys()]):
                    updated_params.append(new_param)

            if len(updated_params) > 0:
                self._update_assay(assay.assay_definition, updated_params)
        except sqlalchemy.exc.StatementError as exc:
            # cancel current transaction (insert)
            self.session.rollback()
            # delete added table
            self.rollback()
            raise OpenLogConnection.ImportException(exc)

    def split_traces(
        self,
        table_name,
        assay_column: AssayColumn,
        discrete: bool,
        geom: str = "proj_effective_geom",
    ):

        assay_table_definition = self.assay_database_definition_map[table_name]

        # split traces for cross section displaying
        queries = self.trace_splitter(
            extent=AssayDataExtent.DISCRETE if discrete else AssayDataExtent.EXTENDED,
            assay_def=assay_table_definition,
            y_col=assay_column.name,
            y_serie_type=assay_column.series_type,
        )

        # queries.get_y_type(self.session)
        queries.set_trace_geometry(geom)

        # if no trace geometry, pass
        if not queries.is_trace_geom_exist(self.session):
            return

        # check if layer exists to avoid multiple call
        if queries.check_if_exists(self.session) and not queries.need_update(
            self.session
        ):
            return

        self.session.execute(queries.remove_if_exists())
        self.session.commit()
        self.session.execute(queries.split_geometry())
        self.session.execute(queries.addgeometrycolumn())
        [self.session.execute(q) for q in queries.updategeometry()]
        self.session.commit()
        # create triggers
        triggers = queries.create_trigger_queries()
        for query in triggers:
            self.session.execute(query)
        self.session.commit()
        # grant permissions
        grants = queries.create_grants_queries()
        for grant in grants:
            self.session.execute(grant)
        self.session.commit()

    @staticmethod
    def _convert_if_numpy_value(value) -> Any:
        """
        Convert numpy in64 or float64 to base type int/float.
        Needed for psycopg2 use with sqlalchemy : https://docs.sqlalchemy.org/en/14/faq/thirdparty.html

        Args:
            value: value to be converted

        Returns: converted value if type is np.int64

        """
        if isinstance(value, np.int64):
            return int(value)
        elif isinstance(value, np.float64):
            return float(value)
        else:
            return value

    def get_assay(self, hole_id: str, variable: str) -> GenericAssay:
        """
        Return GenericAssay for assay data use

        Args:
            hole_id: The collar id of the hole
            variable: variable name of assay

        Returns:
            GenericAssay: generic assay with method to access assay data
        """

        assay_definition_table = (
            self.session.query(self.assay_definition_base)
            .filter(self.assay_definition_base.variable == variable)
            .first()
        )
        if not assay_definition_table:
            raise AssayInterface.ReadException(f"Assay {variable} not available.")
        assay_definition = assay_definition_table.to_assay_definition()

        assay = AssayFactory(
            assay_definition, self.get_assay_table_definition(variable)
        ).create_generic_assay(
            hole_id,
            self.session,
        )

        assay.set_geo_extractor(self.geo_extractor)
        return assay

    def _before_sqlalchemy_table_creation(
        self, table: sqlalchemy.Table, schema: str
    ) -> None:
        """
        Define sqlalchemy events before table creation. Can be used to define specific permission on created table.

        Args:
            table: sqlalchemy table to be created
            schema: (str) database schema used
        """
        pass
