import sqlalchemy
from sqlalchemy import DDL, event
from xplordb.sqlalchemy.assay.structural_types import pgAzimuthType, pgSphericalType

from openlog.core.geo_extractor import XplordbGeoExtractor
from openlog.core.trace_splitter import XplordbSplitTracesQueries
from openlog.datamodel.assay.generic_assay import (
    AssayColumn,
    AssayDatabaseDefinition,
    AssayDataExtent,
    AssayDefinition,
    AssayDomainType,
    AssaySeriesType,
)
from openlog.datamodel.connection.interfaces.categories_interface import (
    CategoriesInterface,
)
from openlog.datamodel.connection.sqlalchemy.sqlalchemy_assay_interface import (
    SqlAlchemyAssayInterface,
)


class XplordbAssayInterface(SqlAlchemyAssayInterface):
    def __init__(self, engine, session, categories_iface: CategoriesInterface):
        """
        Implement AssayInterface for a Xplordb db

        Args:
            engine: sqlalchemy engine
            session: sqlalchemy session created from engine
            categories_iface : CategoriesInterface to get default lith category name
        """
        self._categories_iface = categories_iface
        super().__init__(engine, session, "assay")
        self.trace_splitter = XplordbSplitTracesQueries
        self.geo_extractor = XplordbGeoExtractor

    def default_assay_schema(self) -> str:
        """
        Return default schema for assay table creation

        """
        return "assay"

    def _cast_column_type(
        self, assay_table: sqlalchemy.Table, assay_definition: AssayDefinition
    ) -> sqlalchemy.Table:
        """
        Cast sqlalchemy.Table's columns depending backend engine.
        """
        # define generic function
        def _cast_one_column_type_(
            assay_table: sqlalchemy.Table,
            assay_definition: AssayDefinition,
            source_type: AssaySeriesType,
            target_type: sqlalchemy.types,
        ):
            # some columns have a dedicated postgresql type
            to_remove = [
                n
                for n, v in assay_definition.columns.items()
                if v.series_type == source_type
            ]
            col_instances = [
                col for col in assay_table._columns if col.name in to_remove
            ]
            for col in col_instances:
                assay_table._columns.remove(col)

            # create new columns with specific type
            for col_name in to_remove:
                c = sqlalchemy.Column(col_name, target_type(), primary_key=False)
                assay_table.append_column(c)

            return assay_table

        # spherical : composite type assay.spherical_data
        assay_table = _cast_one_column_type_(
            assay_table, assay_definition, AssaySeriesType.SPHERICAL, pgSphericalType
        )
        # polar : assay.azimuth type
        assay_table = _cast_one_column_type_(
            assay_table, assay_definition, AssaySeriesType.POLAR, pgAzimuthType
        )

        return assay_table

    def _convert_composite_string(
        self, params: list, assay_definition: AssayDefinition
    ) -> list:
        """
        Convert string values to composite type depending backend engine.
        """

        # spherical type has to be converted
        col_names = [
            n
            for n, v in assay_definition.columns.items()
            if v.series_type == AssaySeriesType.SPHERICAL
        ]

        for p in params:
            for name in col_names:
                p[name] = pgSphericalType.insert_from_json(p[name])

        return params

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

        Args:
            base: sqlalchemy base to be created
            schema: (str) database schema used
        """
        schema_str = ""
        if schema:
            schema_str = schema + "."

        event.listen(
            table,
            "after_create",
            DDL(
                f"GRANT SELECT ON TABLE {schema_str}{table.name} TO xdb_viewer;"
                f"GRANT SELECT, UPDATE, INSERT ON TABLE {schema_str}{table.name} TO xdb_logger;"
                f"GRANT SELECT, UPDATE, INSERT ON TABLE {schema_str}{table.name} TO xdb_importer;"
                f"GRANT ALL ON TABLE {schema_str}{table.name} TO xdb_admin;"
            ),
        )
