"""Some helpers to use with Alembic migration tool."""

from alembic.autogenerate import renderers
from alembic.autogenerate import rewriter
from alembic.autogenerate.render import _add_column
from alembic.autogenerate.render import _add_index
from alembic.autogenerate.render import _add_table
from alembic.autogenerate.render import _drop_column
from alembic.autogenerate.render import _drop_index
from alembic.autogenerate.render import _drop_table
from alembic.ddl.base import RenameTable
from alembic.ddl.base import format_table_name
from alembic.ddl.base import visit_rename_table
from alembic.ddl.sqlite import SQLiteImpl
from alembic.operations import BatchOperations
from alembic.operations import Operations
from alembic.operations import ops
from sqlalchemy import Column
from sqlalchemy import text
from sqlalchemy.dialects.mysql.base import MySQLDialect
from sqlalchemy.dialects.sqlite.base import SQLiteDialect
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.schema import DropTable
from sqlalchemy.sql import func
from sqlalchemy.types import TypeDecorator

from geoalchemy2 import Geography
from geoalchemy2 import Geometry
from geoalchemy2 import Raster
from geoalchemy2.admin.dialects.common import _check_spatial_type
from geoalchemy2.admin.dialects.common import _get_gis_cols
from geoalchemy2.admin.dialects.common import _spatial_idx_name

writer = rewriter.Rewriter()
"""Rewriter object for Alembic."""

_SPATIAL_TABLES = set()


class GeoPackageImpl(SQLiteImpl):
    """Class to copy the Alembic implementation from SQLite to GeoPackage."""

    __dialect__ = "geopackage"


def _monkey_patch_get_indexes_for_sqlite():
    """Monkey patch SQLAlchemy to fix spatial index reflection."""
    normal_behavior = SQLiteDialect.get_indexes

    def spatial_behavior(self, connection, table_name, schema=None, **kw):
        indexes = self._get_indexes_normal_behavior(connection, table_name, schema=None, **kw)

        is_gpkg = connection.dialect.name == "geopackage"

        try:
            # Check that SpatiaLite was loaded into the DB
            is_spatial_db = connection.exec_driver_sql(
                """PRAGMA main.table_info({})""".format(
                    "gpkg_geometry_columns" if is_gpkg else "geometry_columns"
                )
            ).fetchall()
            if not is_spatial_db:
                return indexes
        except AttributeError:
            return indexes

        # Get spatial indexes
        if is_gpkg:
            spatial_index_query = text(
                """SELECT A.table_name, A.column_name, IFNULL(B.has_index, 0) AS has_index
                FROM "gpkg_geometry_columns"
                AS A
                LEFT JOIN (
                    SELECT table_name, column_name, COUNT(*) AS has_index
                    FROM gpkg_extensions
                    WHERE LOWER(table_name) = LOWER('{table_name}')
                        AND extension_name = 'gpkg_rtree_index'
                ) AS B
                ON LOWER(A.table_name) = LOWER(B.table_name)
                WHERE LOWER(A.table_name) = LOWER('{table_name}');
            """.format(
                    table_name=table_name
                )
            )
        else:
            spatial_index_query = text(
                """SELECT *
                FROM geometry_columns
                WHERE f_table_name = '{}'
                ORDER BY f_table_name, f_geometry_column;""".format(
                    table_name
                )
            )

        spatial_indexes = connection.execute(spatial_index_query).fetchall()

        if spatial_indexes:
            reflected_names = set([i["name"] for i in indexes])
            for idx in spatial_indexes:
                idx_col = idx[1]
                idx_name = _spatial_idx_name(table_name, idx_col)
                if not bool(idx[-1]) or idx_name in reflected_names:
                    continue
                indexes.append(
                    {
                        "name": idx_name,
                        "column_names": [idx_col],
                        "unique": 0,
                        "dialect_options": {"_column_flag": True},
                    }
                )
                reflected_names.add(idx_name)

        return indexes

    spatial_behavior.__doc__ = normal_behavior.__doc__
    SQLiteDialect.get_indexes = spatial_behavior
    SQLiteDialect._get_indexes_normal_behavior = normal_behavior


_monkey_patch_get_indexes_for_sqlite()


def _monkey_patch_get_indexes_for_mysql():
    """Monkey patch SQLAlchemy to fix spatial index reflection."""
    normal_behavior = MySQLDialect.get_indexes

    def spatial_behavior(self, connection, table_name, schema=None, **kw):
        indexes = self._get_indexes_normal_behavior(connection, table_name, schema=None, **kw)

        # Get spatial indexes
        has_index_query = """SELECT DISTINCT
                COLUMN_NAME
            FROM INFORMATION_SCHEMA.STATISTICS
            WHERE TABLE_NAME = '{}' AND INDEX_TYPE = 'SPATIAL'""".format(
            table_name
        )
        if schema is not None:
            has_index_query += """ AND TABLE_SCHEMA = '{}'""".format(schema)
        spatial_indexes = connection.execute(text(has_index_query)).fetchall()

        if spatial_indexes:
            reflected_names = set([i["name"] for i in indexes])
            for idx in spatial_indexes:
                idx_col = idx[0]
                idx_name = _spatial_idx_name(table_name, idx_col)
                if idx_name in reflected_names:
                    continue
                indexes.append(
                    {
                        "name": idx_name,
                        "column_names": [idx_col],
                        "unique": 0,
                        "dialect_options": {"_column_flag": True},
                    }
                )
                reflected_names.add(idx_name)

        return indexes

    spatial_behavior.__doc__ = normal_behavior.__doc__
    MySQLDialect.get_indexes = spatial_behavior
    MySQLDialect._get_indexes_normal_behavior = normal_behavior


_monkey_patch_get_indexes_for_mysql()


def render_item(obj_type, obj, autogen_context):
    """Add proper imports for spatial types."""
    if obj_type == "type" and isinstance(obj, (Geometry, Geography, Raster)):
        import_name = obj.__class__.__name__
        autogen_context.imports.add(f"from geoalchemy2 import {import_name}")
        return "%r" % obj

    # Default rendering for other objects
    return False


def include_object(obj, name, obj_type, reflected, compare_to):
    """Do not include internal tables of spatial extensions.

    .. warning::
        This function only checks the table names, so it might exclude tables that should not be.
        In such case, you should create your own function to handle your specific table names.

    """
    if obj_type == "table" and (
        name.startswith("geometry_columns")
        or name.startswith("spatial_ref_sys")
        or name.startswith("spatialite_history")
        or name.startswith("sqlite_sequence")
        or name.startswith("views_geometry_columns")
        or name.startswith("virts_geometry_columns")
        or name.startswith("idx_")
        or name.startswith("gpkg_")
        or name.startswith("vgpkg_")
    ):
        return False
    return True


@Operations.register_operation("add_geospatial_column")
@BatchOperations.register_operation("add_geospatial_column", "batch_add_geospatial_column")
class AddGeospatialColumnOp(ops.AddColumnOp):
    """Add a Geospatial Column in an Alembic migration context.

    This method originates from:
    https://alembic.sqlalchemy.org/en/latest/api/operations.html#operation-plugins
    """

    @classmethod
    def add_geospatial_column(cls, operations, table_name, column, schema=None):
        """Handle the different situations arising from adding geospatial column to a DB."""
        op = cls(table_name, column, schema=schema)
        return operations.invoke(op)

    def reverse(self):
        """Used to autogenerate the downgrade function."""
        return DropGeospatialColumnOp.from_column_and_tablename(
            self.schema, self.table_name, self.column.name
        )

    @classmethod
    def batch_add_geospatial_column(
        cls,
        operations,
        column,
        insert_before=None,
        insert_after=None,
    ):
        """Issue an "add column" instruction using the current batch migration context."""
        kw = {}
        if insert_before:
            kw["insert_before"] = insert_before
        if insert_after:
            kw["insert_after"] = insert_after

        op = cls(
            operations.impl.table_name,
            column,
            schema=operations.impl.schema,
            **kw,
        )
        return operations.invoke(op)


@Operations.register_operation("drop_geospatial_column")
@BatchOperations.register_operation("drop_geospatial_column", "batch_drop_geospatial_column")
class DropGeospatialColumnOp(ops.DropColumnOp):
    """Drop a Geospatial Column in an Alembic migration context."""

    @classmethod
    def drop_geospatial_column(cls, operations, table_name, column_name, schema=None, **kw):
        """Handle the different situations arising from dropping geospatial column from a DB."""
        op = cls(table_name, column_name, schema=schema, **kw)
        return operations.invoke(op)

    def reverse(self):
        """Used to autogenerate the downgrade function."""
        return AddGeospatialColumnOp.from_column_and_tablename(
            self.schema, self.table_name, self.column
        )

    @classmethod
    def batch_drop_geospatial_column(cls, operations, column_name, **kw):
        """Issue a "drop column" instruction using the current batch migration context."""
        op = cls(
            operations.impl.table_name,
            column_name,
            schema=operations.impl.schema,
            **kw,
        )
        return operations.invoke(op)


@Operations.implementation_for(AddGeospatialColumnOp)
def add_geospatial_column(operations, operation):
    """Handle the actual column addition according to the dialect backend.

    Args:
        operations: Operations object from alembic base, defining high level migration operations.
        operation: AddGeospatialColumnOp call, with attributes for table_name, column_name,
            column_type, and optional keywords.
    """
    table_name = operation.table_name
    column_name = operation.column.name

    dialect = operations.get_bind().dialect

    if dialect.name == "sqlite":
        if isinstance(operation.column, TypeDecorator):
            # Will be either geoalchemy2.types.Geography or geoalchemy2.types.Geometry, if using a
            # custom type
            geospatial_core_type = operation.column.type.load_dialect_impl(dialect)
        else:
            geospatial_core_type = operation.column.type
        operations.execute(
            func.AddGeometryColumn(
                table_name,
                column_name,
                geospatial_core_type.srid,
                geospatial_core_type.geometry_type,
                geospatial_core_type.dimension,
                not geospatial_core_type.nullable,
            )
        )
    elif "postgresql" in dialect.name:
        operations.impl.add_column(
            table_name,
            operation.column,
            schema=operation.schema,
        )


@Operations.implementation_for(DropGeospatialColumnOp)
def drop_geospatial_column(operations, operation):
    """Handle the actual column removal according to the dialect backend.

    Args:
        operations: Operations object from alembic base, defining high level migration operations.
        operation: AddGeospatialColumnOp call, with attributes for table_name, column_name,
            column_type, and optional keywords.
    """
    table_name = operation.table_name
    column = operation.to_column(operations.migration_context)

    dialect = operations.get_bind().dialect

    if dialect.name == "sqlite":
        _SPATIAL_TABLES.add(table_name)
    operations.impl.drop_column(table_name, column, schema=operation.schema, **operation.kw)


@compiles(RenameTable, "sqlite")
def visit_rename_geospatial_table(element, compiler, **kw):
    """Specific compilation rule to rename spatial tables with SQLite dialect."""
    table_is_spatial = element.table_name in _SPATIAL_TABLES
    new_table_is_spatial = element.new_table_name in _SPATIAL_TABLES

    if table_is_spatial or new_table_is_spatial:
        # Here we suppose there is only one DB attached to the current engine, so the prefix
        # is set to NULL
        return "SELECT RenameTable(NULL, '%s', '%s')" % (
            format_table_name(compiler, element.table_name, element.schema),
            format_table_name(compiler, element.new_table_name, element.schema),
        )
    else:
        return visit_rename_table(element, compiler, **kw)


@compiles(DropTable, "sqlite")
def visit_drop_geospatial_table(element, compiler, **kw):
    """Specific compilation rule to drop spatial tables with SQLite dialect."""
    table_is_spatial = element.element.name in _SPATIAL_TABLES

    if table_is_spatial:
        # Here we suppose there is only one DB attached to the current engine
        return "SELECT DropTable(NULL, '%s')" % (
            format_table_name(compiler, element.element.name, None),
        )
    else:
        return compiler.visit_drop_table(element, **kw)


@renderers.dispatch_for(AddGeospatialColumnOp)
def render_add_geo_column(autogen_context, op):
    """Render the add_geospatial_column operation in migration script."""
    col_render = _add_column(autogen_context, op)
    return col_render.replace(".add_column(", ".add_geospatial_column(")


@renderers.dispatch_for(DropGeospatialColumnOp)
def render_drop_geo_column(autogen_context, op):
    """Render the drop_geospatial_column operation in migration script."""
    col_render = _drop_column(autogen_context, op)
    return col_render.replace(".drop_column(", ".drop_geospatial_column(")


@writer.rewrites(ops.AddColumnOp)
def add_geo_column(context, revision, op):
    """Replace the default AddColumnOp by a geospatial-specific one."""
    col_type = op.column.type
    if isinstance(col_type, TypeDecorator):
        dialect = context.bind.dialect
        col_type = col_type.load_dialect_impl(dialect)
    if isinstance(col_type, (Geometry, Geography, Raster)):
        op.column.type.spatial_index = False
        op.column.type._spatial_index_reflected = None
        new_op = AddGeospatialColumnOp(op.table_name, op.column, schema=op.schema)
    else:
        new_op = op
    return new_op


@writer.rewrites(ops.DropColumnOp)
def drop_geo_column(context, revision, op):
    """Replace the default DropColumnOp by a geospatial-specific one."""
    col_type = op.to_column().type
    if isinstance(col_type, TypeDecorator):
        dialect = context.bind.dialect
        col_type = col_type.load_dialect_impl(dialect)
    if isinstance(col_type, (Geometry, Geography, Raster)):
        new_op = DropGeospatialColumnOp(op.table_name, op.column_name, schema=op.schema)
    else:
        new_op = op
    return new_op


@Operations.register_operation("create_geospatial_table")
class CreateGeospatialTableOp(ops.CreateTableOp):
    """Create a Geospatial Table in an Alembic migration context.

    This method originates from:
    https://alembic.sqlalchemy.org/en/latest/api/operations.html#operation-plugins
    """

    @classmethod
    def create_geospatial_table(cls, operations, table_name, *columns, **kw):
        """Handle the different situations arising from creating geospatial table to a DB."""
        op = cls(table_name, columns, **kw)
        return operations.invoke(op)

    def reverse(self):
        """Used to autogenerate the downgrade function."""
        return DropGeospatialColumnOp.from_table(
            self.to_table(),
            _namespace_metadata=self._namespace_metadata,
        )

    @classmethod
    def from_table(cls, table, _namespace_metadata=None):
        obj = super().from_table(table, _namespace_metadata)
        return obj

    def to_table(self, migration_context=None):
        table = super().to_table(migration_context)

        # Set spatial_index attribute to False so the indexes are created explicitly
        for col in table.columns:
            try:
                if col.type.spatial_index:
                    col.type.spatial_index = False
            except AttributeError:
                pass
        return table


@Operations.register_operation("drop_geospatial_table")
class DropGeospatialTableOp(ops.DropTableOp):
    @classmethod
    def drop_geospatial_table(cls, operations, table_name, schema=None, **kw):
        """Handle the different situations arising from dropping geospatial table from a DB."""
        op = cls(table_name, schema=schema, table_kw=kw)
        return operations.invoke(op)

    def reverse(self):
        """Used to autogenerate the downgrade function."""
        return CreateGeospatialTableOp.from_table(
            self.to_table(),
            _namespace_metadata=self._namespace_metadata,
        )

    @classmethod
    def from_table(cls, table, _namespace_metadata=None):
        obj = super().from_table(table, _namespace_metadata)
        return obj

    def to_table(self, migration_context=None):
        table = super().to_table(migration_context)
        return table


@Operations.implementation_for(CreateGeospatialTableOp)
def create_geospatial_table(operations, operation):
    """Handle the actual table creation according to the dialect backend.

    Args:
        operations: Operations object from alembic base, defining high level migration operations.
        operation: CreateGeospatialTableOp call, with attributes for table_name, column_name,
            column_type, and optional keywords.
    """
    table_name = operation.table_name
    bind = operations.get_bind()

    # For now the default events defined in geoalchemy2 are enough to handle table creation
    operations.create_table(table_name, *operation.columns, schema=operation.schema, **operation.kw)

    if bind.dialect.name == "sqlite":
        _SPATIAL_TABLES.add(table_name)


@Operations.implementation_for(DropGeospatialTableOp)
def drop_geospatial_table(operations, operation):
    """Handle the actual table removal according to the dialect backend.

    Args:
        operations: Operations object from alembic base, defining high level migration operations.
        operation: DropGeospatialTableOp call, with attributes for table_name, column_name,
            column_type, and optional keywords.
    """
    table_name = operation.table_name
    bind = operations.get_bind()
    dialect = bind.dialect

    if dialect.name == "sqlite":
        _SPATIAL_TABLES.add(table_name)
    operations.drop_table(table_name, schema=operation.schema, **operation.table_kw)


@renderers.dispatch_for(CreateGeospatialTableOp)
def render_create_geo_table(autogen_context, op):
    """Render the create_geospatial_table operation in migration script."""
    table_render = _add_table(autogen_context, op)
    return table_render.replace(".create_table(", ".create_geospatial_table(")


@renderers.dispatch_for(DropGeospatialTableOp)
def render_drop_geo_table(autogen_context, op):
    """Render the drop_geospatial_table operation in migration script."""
    table_render = _drop_table(autogen_context, op)
    return table_render.replace(".drop_table(", ".drop_geospatial_table(")


@writer.rewrites(ops.CreateTableOp)
def create_geo_table(context, revision, op):
    """Replace the default CreateTableOp by a geospatial-specific one."""
    dialect = context.bind.dialect
    gis_cols = _get_gis_cols(op, (Geometry, Geography, Raster), dialect)

    if gis_cols:
        new_op = CreateGeospatialTableOp(
            op.table_name,
            op.columns,
            schema=op.schema,
            _namespace_metadata=op._namespace_metadata,
            _constraints_included=op._constraints_included,
            **op.kw,
        )
    else:
        new_op = op

    return new_op


@writer.rewrites(ops.DropTableOp)
def drop_geo_table(context, revision, op):
    """Replace the default DropTableOp by a geospatial-specific one."""
    dialect = context.bind.dialect
    table = op.to_table()
    gis_cols = _get_gis_cols(table, (Geometry, Geography, Raster), dialect)

    if gis_cols:
        new_op = DropGeospatialTableOp(op.table_name, schema=op.schema)
    else:
        new_op = op

    return new_op


@Operations.register_operation("create_geospatial_index")
@BatchOperations.register_operation("create_geospatial_index", "batch_create_geospatial_index")
class CreateGeospatialIndexOp(ops.CreateIndexOp):
    @classmethod
    def create_geospatial_index(
        cls,
        operations,
        index_name,
        table_name,
        columns,
        schema=None,
        unique=False,
        **kw,
    ):
        """Handle the different situations arising from creating geospatial index into a DB."""
        op = cls(index_name, table_name, columns, schema=schema, unique=unique, **kw)
        return operations.invoke(op)

    def reverse(self):
        """Used to autogenerate the downgrade function."""
        return DropGeospatialIndexOp(
            self.index_name,
            self.table_name,
            column_name=self.columns[0].name,
            schema=self.schema,
        )

    @classmethod
    def batch_create_geospatial_index(
        cls,
        operations,
        index_name,
        columns,
        **kw,
    ):
        """Issue a "create index" instruction using the current batch migration context."""
        op = cls(
            index_name,
            operations.impl.table_name,
            columns,
            schema=operations.impl.schema,
            **kw,
        )
        return operations.invoke(op)


@Operations.register_operation("drop_geospatial_index")
@BatchOperations.register_operation("drop_geospatial_index", "batch_drop_geospatial_index")
class DropGeospatialIndexOp(ops.DropIndexOp):
    def __init__(self, *args, column_name, **kwargs):
        super().__init__(*args, **kwargs)
        self.column_name = column_name

    @classmethod
    def drop_geospatial_index(
        cls,
        operations,
        index_name,
        table_name,
        column_name,
        schema=None,
        unique=False,
        **kw,
    ):
        """Handle the different situations arising from dropping geospatial index from a DB."""
        op = cls(
            index_name,
            table_name=table_name,
            column_name=column_name,
            schema=schema,
            unique=unique,
            **kw,
        )
        return operations.invoke(op)

    def reverse(self):
        """Used to autogenerate the downgrade function."""
        return CreateGeospatialIndexOp(
            self.index_name,
            self.table_name,
            column_name=self.column_name,
            schema=self.schema,
            _reverse=self,
            **self.kw,
        )

    @classmethod
    def from_index(cls, index):
        assert index.table is not None
        assert len(index.columns) == 1, "A spatial index must be set on one column only"
        return cls(
            index.name,
            index.table.name,
            column_name=index.columns[0].name,
            schema=index.table.schema,
            _reverse=CreateGeospatialIndexOp.from_index(index),
            **index.kwargs,
        )

    @classmethod
    def batch_drop_geospatial_index(cls, operations, index_name, **kw):
        """Issue a "drop index" instruction using the current batch migration context."""
        op = cls(
            index_name,
            table_name=operations.impl.table_name,
            schema=operations.impl.schema,
            **kw,
        )
        return operations.invoke(op)


@Operations.implementation_for(CreateGeospatialIndexOp)
def create_geospatial_index(operations, operation):
    """Handle the actual index creation according to the dialect backend.

    Args:
        operations: Operations object from alembic base, defining high level migration operations.
        operation: CreateGeospatialIndexOp call, with attributes for table_name, column_name,
            column_type, and optional keywords.
    """
    bind = operations.get_bind()
    dialect = bind.dialect

    if dialect.name == "sqlite":
        assert len(operation.columns) == 1, "A spatial index must be set on one column only"
        operations.execute(func.CreateSpatialIndex(operation.table_name, operation.columns[0]))
    else:
        idx = operation.to_index(operations.migration_context)
        operations.impl.create_index(idx)


@Operations.implementation_for(DropGeospatialIndexOp)
def drop_geospatial_index(operations, operation):
    """Handle the actual index drop according to the dialect backend.

    Args:
        operations: Operations object from alembic base, defining high level migration operations.
        operation: DropGeospatialIndexOp call, with attributes for table_name, column_name,
            column_type, and optional keywords.
    """
    bind = operations.get_bind()
    dialect = bind.dialect

    if dialect.name == "sqlite":
        operations.execute(func.DisableSpatialIndex(operation.table_name, operation.column_name))
    else:
        operations.impl.drop_index(operation.to_index(operations.migration_context))


@renderers.dispatch_for(CreateGeospatialIndexOp)
def render_create_geo_index(autogen_context, op):
    """Render the create_geospatial_index operation in migration script."""
    idx_render = _add_index(autogen_context, op)
    return idx_render.replace(".create_index(", ".create_geospatial_index(")


@renderers.dispatch_for(DropGeospatialIndexOp)
def render_drop_geo_index(autogen_context, op):
    """Render the drop_geospatial_index operation in migration script."""
    idx_render = _drop_index(autogen_context, op)

    # Replace function name
    text = idx_render.replace(".drop_index(", ".drop_geospatial_index(")

    # Add column name as keyword argument
    text = text[:-1] + ", column_name='%s')" % (op.column_name,)

    return text


@writer.rewrites(ops.CreateIndexOp)
def create_geo_index(context, revision, op):
    """Replace the default CreateIndexOp by a geospatial-specific one."""
    dialect = context.bind.dialect

    if len(op.columns) == 1:
        col = op.columns[0]
        if isinstance(col, Column) and _check_spatial_type(
            col.type, (Geometry, Geography, Raster), dialect
        ):
            # Fix index properties
            op.kw["postgresql_using"] = op.kw.get("postgresql_using", "gist")
            if col.type.use_N_D_index:
                postgresql_ops = {col.name: "gist_geometry_ops_nd"}
            else:
                postgresql_ops = {}
            op.kw["postgresql_ops"] = op.kw.get("postgresql_ops", postgresql_ops)

            return CreateGeospatialIndexOp(
                op.index_name,
                op.table_name,
                op.columns,
                schema=op.schema,
                unique=op.unique,
                **op.kw,
            )

    return op


@writer.rewrites(ops.DropIndexOp)
def drop_geo_index(context, revision, op):
    """Replace the default DropIndexOp by a geospatial-specific one."""
    dialect = context.bind.dialect
    idx = op.to_index()

    if len(idx.columns) == 1:
        col = idx.columns[0]
        if isinstance(col, Column) and _check_spatial_type(
            col.type, (Geometry, Geography, Raster), dialect
        ):
            return DropGeospatialIndexOp(
                op.index_name,
                table_name=op.table_name,
                column_name=col.name,
                schema=op.schema,
                **op.kw,
            )

    return op