from typing import Any, List

import numpy as np
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, PlgTranslator


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.tr = PlgTranslator().tr
        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
        self.assay_database_definition_map = self._init_assay_maps()

        # 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=self.tr(
                    "Exception raised for assay tables creation. It must come from already created "
                    "tables : {}."
                ).format(exc)
            )

    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_tables = self.session.query(self.assay_definition_base).all()
        assay_def = [assay_table.to_assay_definition() for assay_table in assay_tables]

        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
        """
        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]) -> 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)

        # create a layer for each column
        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
                # effective 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="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="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) -> None:
        """
        Import assay into OpenLogConnection

        Args:
            assay: (GenericAssay) assay to be imported
        """
        # 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}

            if assay.assay_definition.data_extent == AssayDataExtent.DISCRETE:
                value[assay_table_definition.x_col] = self._convert_if_numpy_int(
                    assay.import_x_values[i]
                )
            else:
                value[assay_table_definition.x_col] = self._convert_if_numpy_int(
                    assay.import_x_values[i][0]
                )
                value[assay_table_definition.x_end_col] = self._convert_if_numpy_int(
                    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_int(vals[i])

            params.append(value)

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

        try:
            self.session.execute(assay_table.insert(), 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 = "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.set_trace_geometry(geom)

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

        self.session.execute(queries.remove_if_exists())
        self.session.commit()
        queries.get_y_type(self.session)
        self.session.execute(queries.split_geometry())
        self.session.execute(queries.addgeometrycolumn())
        self.session.execute(queries.updategeometry())
        self.session.commit()
        # create triggers
        triggers = queries.create_trigger_queries()
        for query in triggers:
            self.session.execute(query)
        self.session.commit()

    @staticmethod
    def _convert_if_numpy_int(value) -> Any:
        """
        Convert numpy in64 to base type int.
        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)
        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
