"""This module defines specific functions for SQLite dialect."""

import os
from typing import Optional

from sqlalchemy import text
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql import func
from sqlalchemy.sql import select

from geoalchemy2 import functions
from geoalchemy2.admin.dialects.common import _check_spatial_type
from geoalchemy2.admin.dialects.common import _format_select_args
from geoalchemy2.admin.dialects.common import _spatial_idx_name
from geoalchemy2.admin.dialects.common import setup_create_drop
from geoalchemy2.types import Geography
from geoalchemy2.types import Geometry
from geoalchemy2.types import _DummyGeometry
from geoalchemy2.utils import authorized_values_in_docstring


def load_spatialite_driver(dbapi_conn, *args):
    """Load SpatiaLite extension in SQLite connection.

    .. Warning::
        The path to the SpatiaLite module should be set in the `SPATIALITE_LIBRARY_PATH`
        environment variable.

    Args:
        dbapi_conn: The DBAPI connection.
    """
    if "SPATIALITE_LIBRARY_PATH" not in os.environ:
        raise RuntimeError("The SPATIALITE_LIBRARY_PATH environment variable is not set.")
    dbapi_conn.enable_load_extension(True)
    dbapi_conn.load_extension(os.environ["SPATIALITE_LIBRARY_PATH"])
    dbapi_conn.enable_load_extension(False)


_JOURNAL_MODE_VALUES = ["DELETE", "TRUNCATE", "PERSIST", "MEMORY", "WAL", "OFF"]


@authorized_values_in_docstring(JOURNAL_MODE_VALUES=_JOURNAL_MODE_VALUES)
def init_spatialite(
    dbapi_conn,
    *args,
    transaction: bool = False,
    init_mode: Optional[str] = None,
    journal_mode: Optional[str] = None,
):
    """Initialize internal SpatiaLite tables.

    Args:
        dbapi_conn: The DBAPI connection.
        transaction: If set to `True` the whole operation will be handled as a single Transaction
            (faster). The default value is `False` (slower, but safer).
        init_mode: Can be `None` to load all EPSG SRIDs, `'WGS84'` to load only the ones related
            to WGS84 or `'EMPTY'` to not load any EPSG SRID.

            .. Note::

                It is possible to load other EPSG SRIDs afterwards using `InsertEpsgSrid(srid)`.

        journal_mode: Change the journal mode to the given value. This can make the table creation
            much faster. The possible values are the following: <JOURNAL_MODE_VALUES>. See
            https://www.sqlite.org/pragma.html#pragma_journal_mode for more details.

            .. Warning::
                Some values, like 'MEMORY' or 'OFF', can lead to corrupted databases if the process
                is interrupted during initialization.

            .. Note::
                The original value is restored after the initialization.

    .. Note::
        When using this function as a listener it is not possible to pass the `transaction`,
        `init_mode` or `journal_mode` arguments directly. To do this you can either create another
        function that calls `init_spatialite` (or
        :func:`geoalchemy2.admin.dialects.sqlite.load_spatialite` if you also want to load the
        SpatiaLite drivers) with an hard-coded `init_mode` or just use a lambda::

            >>> sqlalchemy.event.listen(
            ...     engine,
            ...     "connect",
            ...     lambda x, y: init_spatialite(
            ...         x,
            ...         y,
            ...         transaction=True,
            ...         init_mode="EMPTY",
            ...         journal_mode="OFF",
            ...     )
            ... )
    """
    func_args = []

    # Check the value of the 'transaction' parameter
    if not isinstance(transaction, (bool, int)):
        raise ValueError("The 'transaction' argument must be True or False.")
    else:
        func_args.append(str(transaction))

    # Check the value of the 'init_mode' parameter
    init_mode_values = ["WGS84", "EMPTY"]
    if isinstance(init_mode, str):
        init_mode = init_mode.upper()
    if init_mode is not None:
        if init_mode not in init_mode_values:
            raise ValueError("The 'init_mode' argument must be one of {}.".format(init_mode_values))
        func_args.append(f"'{init_mode}'")

    # Check the value of the 'journal_mode' parameter
    if isinstance(journal_mode, str):
        journal_mode = journal_mode.upper()
    if journal_mode is not None:
        if journal_mode not in _JOURNAL_MODE_VALUES:
            raise ValueError(
                "The 'journal_mode' argument must be one of {}.".format(_JOURNAL_MODE_VALUES)
            )

    if dbapi_conn.execute("SELECT CheckSpatialMetaData();").fetchone()[0] < 1:
        if journal_mode is not None:
            current_journal_mode = dbapi_conn.execute("PRAGMA journal_mode").fetchone()[0]
            dbapi_conn.execute("PRAGMA journal_mode = {}".format(journal_mode))

        dbapi_conn.execute("SELECT InitSpatialMetaData({});".format(", ".join(func_args)))

        if journal_mode is not None:
            dbapi_conn.execute("PRAGMA journal_mode = {}".format(current_journal_mode))


def load_spatialite(dbapi_conn, *args, **kwargs):
    """Load SpatiaLite extension in SQLite DB and initialize internal tables.

    See :func:`geoalchemy2.admin.dialects.sqlite.load_spatialite_driver` and
    :func:`geoalchemy2.admin.dialects.sqlite.init_spatialite` functions for details about
    arguments.
    """
    load_spatialite_driver(dbapi_conn)
    init_spatialite(dbapi_conn, **kwargs)


def _get_spatialite_attrs(bind, table_name, col_name):
    attrs = bind.execute(
        text(
            """SELECT * FROM "geometry_columns"
            WHERE LOWER(f_table_name) = LOWER(:table_name)
                AND LOWER(f_geometry_column) = LOWER(:column_name)
        """
        ).bindparams(table_name=table_name, column_name=col_name)
    ).fetchone()
    if attrs is None:
        # If the column is not registered as a spatial column we ignore it
        return None
    return attrs[2:]


def get_spatialite_version(bind):
    """Get the version of the currently loaded Spatialite extension."""
    return bind.execute(text("SELECT spatialite_version();")).fetchone()[0]


def _setup_dummy_type(table, gis_cols):
    """Setup dummy type for new Geometry columns so they can be updated later."""
    for col in gis_cols:
        # Add dummy columns with GEOMETRY type
        col._actual_type = col.type
        col.type = _DummyGeometry()
    table.columns = table.info["_saved_columns"]


def get_col_dim(col):
    """Get dimension of the column type."""
    if col.type.dimension == 4:
        dimension = "XYZM"
    elif col.type.dimension == 2:
        dimension = "XY"
    else:
        if col.type.geometry_type.endswith("M"):
            dimension = "XYM"
        else:
            dimension = "XYZ"
    return dimension


def create_spatial_index(bind, table, col):
    """Create spatial index on the given column."""
    stmt = select(*_format_select_args(func.CreateSpatialIndex(table.name, col.name)))
    stmt = stmt.execution_options(autocommit=True)
    bind.execute(stmt)


def disable_spatial_index(bind, table, col):
    """Disable spatial indexes if present."""
    stmt = select(*_format_select_args(func.CheckSpatialIndex(table.name, col.name)))
    if bind.execute(stmt).fetchone()[0] is not None:
        stmt = select(*_format_select_args(func.DisableSpatialIndex(table.name, col.name)))
        stmt = stmt.execution_options(autocommit=True)
        bind.execute(stmt)
        bind.execute(
            text(
                "DROP TABLE IF EXISTS {};".format(
                    _spatial_idx_name(
                        table.name,
                        col.name,
                    )
                )
            )
        )


def reflect_geometry_column(inspector, table, column_info):
    """Reflect a column of type Geometry with SQLite dialect."""
    # Get geometry type, SRID and spatial index from the SpatiaLite metadata
    if not isinstance(column_info.get("type"), Geometry):
        return
    col_attributes = _get_spatialite_attrs(inspector.bind, table.name, column_info["name"])
    if col_attributes is not None:
        geometry_type, coord_dimension, srid, spatial_index = col_attributes

        if isinstance(geometry_type, int):
            geometry_type_str = str(geometry_type)
            if geometry_type >= 1000:
                first_digit = geometry_type_str[0]
                has_z = first_digit in ["1", "3"]
                has_m = first_digit in ["2", "3"]
            else:
                has_z = has_m = False
            geometry_type = {
                "0": "GEOMETRY",
                "1": "POINT",
                "2": "LINESTRING",
                "3": "POLYGON",
                "4": "MULTIPOINT",
                "5": "MULTILINESTRING",
                "6": "MULTIPOLYGON",
                "7": "GEOMETRYCOLLECTION",
            }[geometry_type_str[-1]]
            if has_z:
                geometry_type += "Z"
            if has_m:
                geometry_type += "M"
        else:
            if "Z" in coord_dimension and "Z" not in geometry_type[-2:]:
                geometry_type += "Z"
            if "M" in coord_dimension and "M" not in geometry_type[-2:]:
                geometry_type += "M"
            coord_dimension = {
                "XY": 2,
                "XYZ": 3,
                "XYM": 3,
                "XYZM": 4,
            }.get(coord_dimension, coord_dimension)

        # Set attributes
        column_info["type"].geometry_type = geometry_type
        column_info["type"].dimension = coord_dimension
        column_info["type"].srid = srid
        column_info["type"].spatial_index = bool(spatial_index)

        # Spatial indexes are not automatically reflected with SQLite dialect
        column_info["type"]._spatial_index_reflected = False


def connect(dbapi_conn, *args, **kwargs):
    """Even handler to load spatial extension when a new connection is created."""
    return load_spatialite(dbapi_conn, *args, **kwargs)


def before_create(table, bind, **kw):
    """Handle spatial indexes during the before_create event."""
    dialect, gis_cols, regular_cols = setup_create_drop(table, bind)

    # Remove the spatial indexes from the table metadata because they should not be
    # created during the table.create() step since the associated columns do not exist
    # at this time.
    table.info["_after_create_indexes"] = []
    current_indexes = set(table.indexes)
    for idx in current_indexes:
        for col in table.info["_saved_columns"]:
            if (_check_spatial_type(col.type, Geometry, dialect)) and col in idx.columns.values():
                table.indexes.remove(idx)
                if idx.name != _spatial_idx_name(table.name, col.name) or not getattr(
                    col.type, "spatial_index", False
                ):
                    table.info["_after_create_indexes"].append(idx)

    _setup_dummy_type(table, gis_cols)


def after_create(table, bind, **kw):
    """Handle spatial indexes during the after_create event."""
    dialect = bind.dialect

    table.columns = table.info.pop("_saved_columns")
    for col in table.columns:
        # Add the managed Geometry columns with RecoverGeometryColumn()
        if _check_spatial_type(col.type, Geometry, dialect):
            col.type = col._actual_type
            del col._actual_type
            dimension = get_col_dim(col)
            args = [table.name, col.name, col.type.srid, col.type.geometry_type, dimension]

            stmt = select(*_format_select_args(func.RecoverGeometryColumn(*args)))
            stmt = stmt.execution_options(autocommit=True)
            bind.execute(stmt)

    for col in table.columns:
        # Add spatial indexes for the Geometry and Geography columns
        # TODO: Check that the Geography type makes sense here
        if (
            _check_spatial_type(col.type, (Geometry, Geography), dialect)
            and col.type.spatial_index is True
        ):
            create_spatial_index(bind, table, col)

    for idx in table.info.pop("_after_create_indexes"):
        table.indexes.add(idx)
        idx.create(bind=bind)


def before_drop(table, bind, **kw):
    """Handle spatial indexes during the before_drop event."""
    dialect, gis_cols, regular_cols = setup_create_drop(table, bind)

    for col in gis_cols:
        # Disable spatial indexes if present
        disable_spatial_index(bind, table, col)

        args = [table.name, col.name]

        stmt = select(*_format_select_args(func.DiscardGeometryColumn(*args)))
        stmt = stmt.execution_options(autocommit=True)
        bind.execute(stmt)


def after_drop(table, bind, **kw):
    """Handle spatial indexes during the after_drop event."""
    table.columns = table.info.pop("_saved_columns")


# Define compiled versions for functions in SpatiaLite whose names don't have
# the ST_ prefix.
_SQLITE_FUNCTIONS = {
    "ST_GeomFromEWKT": "GeomFromEWKT",
    "ST_GeomFromEWKB": "GeomFromEWKB",
    "ST_AsBinary": "AsBinary",
    "ST_AsEWKB": "AsEWKB",
    "ST_AsGeoJSON": "AsGeoJSON",
}


def _compiles_sqlite(cls, fn):
    def _compile_sqlite(element, compiler, **kw):
        return "{}({})".format(fn, compiler.process(element.clauses, **kw))

    compiles(getattr(functions, cls), "sqlite")(_compile_sqlite)


def register_sqlite_mapping(mapping):
    """Register compilation mappings for the given functions.

    Args:
        mapping: Should have the following form::

                {
                    "function_name_1": "sqlite_function_name_1",
                    "function_name_2": "sqlite_function_name_2",
                    ...
                }
    """
    for cls, fn in mapping.items():
        _compiles_sqlite(cls, fn)


register_sqlite_mapping(_SQLITE_FUNCTIONS)
