"""This module defines the Column types.

The :class:`geoalchemy2.types.Geometry`, :class:`geoalchemy2.types.Geography`, and
:class:`geoalchemy2.types.Raster` classes are used when defining geometry, geography and raster
columns/properties in models.
"""
import warnings

from sqlalchemy.dialects import postgresql
from sqlalchemy.dialects.postgresql.base import ischema_names as postgresql_ischema_names
from sqlalchemy.dialects.sqlite.base import ischema_names as sqlite_ischema_names
from sqlalchemy.sql import func
from sqlalchemy.types import Float
from sqlalchemy.types import Integer
from sqlalchemy.types import UserDefinedType

try:
    from .shape import to_shape

    SHAPELY = True
except ImportError:
    SHAPELY = False


from .comparator import BaseComparator
from .comparator import Comparator
from .elements import CompositeElement
from .elements import RasterElement
from .elements import WKBElement
from .elements import WKTElement
from .exc import ArgumentError


class _GISType(UserDefinedType):
    """The base class for spatial types.

    This class defines ``bind_expression`` and ``column_expression`` methods
    that wrap column expressions in ``ST_GeomFromEWKT``, ``ST_GeogFromText``,
    or ``ST_AsEWKB`` calls.

    This class also defines ``result_processor`` and ``bind_processor``
    methods. The function returned by ``result_processor`` converts WKB values
    received from the database to :class:`geoalchemy2.elements.WKBElement`
    objects. The function returned by ``bind_processor`` converts
    :class:`geoalchemy2.elements.WKTElement` objects to EWKT strings.

    Args:
        geometry_type: The geometry type.

            Possible values are:

              * ``"GEOMETRY"``,
              * ``"POINT"``,
              * ``"LINESTRING"``,
              * ``"POLYGON"``,
              * ``"MULTIPOINT"``,
              * ``"MULTILINESTRING"``,
              * ``"MULTIPOLYGON"``,
              * ``"GEOMETRYCOLLECTION"``,
              * ``"CURVE"``,
              * ``None``.

           The latter is actually not supported with
           :class:`geoalchemy2.types.Geography`.

           When set to ``None`` then no "geometry type" constraints will be
           attached to the geometry type declaration. Using ``None`` here
           is not compatible with setting ``management`` to ``True``.

           Default is ``"GEOMETRY"``.

        srid: The SRID for this column. E.g. 4326. Default is ``-1``.
        dimension: The dimension of the geometry. Default is ``2``.

            With ``management`` set to ``True``, that is when ``AddGeometryColumn`` is used
            to add the geometry column, there are two constraints:

            * The ``geometry_type`` must not end with ``"ZM"``.  This is due to PostGIS'
              ``AddGeometryColumn`` failing with ZM geometry types. Instead the "simple"
              geometry type (e.g. POINT rather POINTZM) should be used with ``dimension``
              set to ``4``.
            * When the ``geometry_type`` ends with ``"Z"`` or ``"M"`` then ``dimension``
              must be set to ``3``.

            With ``management`` set to ``False`` (the default) ``dimension`` is not
            taken into account, and the actual dimension is fully defined with the
            ``geometry_type``.

        spatial_index: Indicate if a spatial index should be created. Default is ``True``.
        use_N_D_index: Use the N-D index instead of the standard 2-D index.
        management: Indicate if the ``AddGeometryColumn`` and ``DropGeometryColumn``
            managements functions should be called when adding and dropping the
            geometry column. Should be set to ``True`` for PostGIS 1.x and SQLite. Default is
            ``False``. Note that this option has no effect for
            :class:`geoalchemy2.types.Geography`.
        use_typmod: By default PostgreSQL type modifiers are used to create the geometry
            column. To use check constraints instead set ``use_typmod`` to
            ``False``. By default this option is not included in the call to
            ``AddGeometryColumn``. Note that this option is only taken
            into account if ``management`` is set to ``True`` and is only available
            for PostGIS 2.x.
    """

    name = None
    """ Name used for defining the main geo type (geometry or geography)
        in CREATE TABLE statements. Set in subclasses. """

    from_text = None
    """ The name of "from text" function for this type.
        Set in subclasses. """

    as_binary = None
    """ The name of the "as binary" function for this type.
        Set in subclasses. """

    comparator_factory = Comparator
    """ This is the way by which spatial operators are defined for
        geometry/geography columns. """

    cache_ok = False
    """ Disable cache for this type. """

    def __init__(
        self,
        geometry_type="GEOMETRY",
        srid=-1,
        dimension=2,
        spatial_index=True,
        use_N_D_index=False,
        management=False,
        use_typmod=None,
        from_text=None,
        name=None,
        nullable=True,
        _spatial_index_reflected=None,
    ):
        geometry_type, srid = self.check_ctor_args(
            geometry_type, srid, dimension, management, use_typmod, nullable
        )
        self.geometry_type = geometry_type
        self.srid = srid
        if name is not None:
            self.name = name
        if from_text is not None:
            self.from_text = from_text
        self.dimension = dimension
        self.spatial_index = spatial_index
        self.use_N_D_index = use_N_D_index
        if management:
            warnings.warn(
                "The 'management' parameter is going to be deprecated and will raise an error in "
                "the version 0.14",
                PendingDeprecationWarning,
            )
        self.management = management
        self.use_typmod = use_typmod
        self.extended = self.as_binary == "ST_AsEWKB"
        self.nullable = nullable
        self._spatial_index_reflected = _spatial_index_reflected

    def get_col_spec(self):
        if not self.geometry_type:
            return self.name
        return "%s(%s,%d)" % (self.name, self.geometry_type, self.srid)

    def column_expression(self, col):
        """Specific column_expression that automatically adds a conversion function."""
        return getattr(func, self.as_binary)(col, type_=self)

    def result_processor(self, dialect, coltype):
        """Specific result_processor that automatically process spatial elements."""

        def process(value):
            if value is not None:
                kwargs = {}
                if self.srid > 0:
                    kwargs["srid"] = self.srid
                if self.extended is not None:
                    kwargs["extended"] = self.extended
                return self.ElementType(value, **kwargs)

        return process

    def bind_expression(self, bindvalue):
        """Specific bind_expression that automatically adds a conversion function."""
        return getattr(func, self.from_text)(bindvalue, type_=self)

    def bind_processor(self, dialect):
        """Specific bind_processor that automatically process spatial elements."""

        def process(bindvalue):
            if isinstance(bindvalue, WKTElement):
                if bindvalue.extended:
                    return "%s" % (bindvalue.data)
                else:
                    return "SRID=%d;%s" % (bindvalue.srid, bindvalue.data)
            elif isinstance(bindvalue, WKBElement):
                if dialect.name == "sqlite" or not bindvalue.extended:
                    # With SpatiaLite or when the WKBElement includes a WKB value rather
                    # than a EWKB value we use Shapely to convert the WKBElement to an
                    # EWKT string
                    if not SHAPELY:
                        raise ArgumentError(
                            "Shapely is required for handling WKBElement bind "
                            "values when using SpatiaLite or when the bind value "
                            "is a WKB rather than an EWKB"
                        )
                    shape = to_shape(bindvalue)
                    return "SRID=%d;%s" % (bindvalue.srid, shape.wkt)
                else:
                    # PostGIS ST_GeomFromEWKT works with EWKT strings as well
                    # as EWKB hex strings
                    return bindvalue.desc
            elif isinstance(bindvalue, RasterElement):
                return "%s" % (bindvalue.data)
            else:
                return bindvalue

        return process

    @staticmethod
    def check_ctor_args(geometry_type, srid, dimension, management, use_typmod, nullable):
        try:
            srid = int(srid)
        except ValueError:
            raise ArgumentError("srid must be convertible to an integer")
        if geometry_type:
            geometry_type = geometry_type.upper()
            if management:
                if geometry_type.endswith("ZM"):
                    # PostGIS' AddGeometryColumn does not work with ZM geometry types. Instead
                    # the simple geometry type (e.g. POINT rather POINTZM) should be used with
                    # dimension set to 4
                    raise ArgumentError(
                        "with management=True use geometry_type={!r} and "
                        "dimension=4 for {!r} geometries".format(geometry_type[:-2], geometry_type)
                    )
                elif geometry_type[-1] in ("Z", "M") and dimension != 3:
                    # If a Z or M geometry type is used then dimension must be set to 3
                    raise ArgumentError(
                        "with management=True dimension must be 3 for "
                        "{!r} geometries".format(geometry_type)
                    )
        else:
            if management:
                raise ArgumentError("geometry_type set to None not compatible with management")
            if srid > 0:
                warnings.warn("srid not enforced when geometry_type is None")

        if use_typmod and not management:
            warnings.warn("use_typmod ignored when management is False")
        if use_typmod is not None and not nullable:
            raise ArgumentError(
                'The "nullable" and "use_typmod" arguments can not be used together'
            )

        return geometry_type, srid


class Geometry(_GISType):
    """The Geometry type.

    Creating a geometry column is done like this::

        Column(Geometry(geometry_type='POINT', srid=4326))

    See :class:`geoalchemy2.types._GISType` for the list of arguments that can
    be passed to the constructor.

    If ``srid`` is set then the ``WKBElement`` objects resulting from queries will
    have that SRID, and, when constructing the ``WKBElement`` objects, the SRID
    won't be read from the data returned by the database. If ``srid`` is not set
    (meaning it's ``-1``) then the SRID set in ``WKBElement`` objects will be read
    from the data returned by the database.
    """

    name = "geometry"
    """ Type name used for defining geometry columns in ``CREATE TABLE``. """

    from_text = "ST_GeomFromEWKT"
    """ The "from text" geometry constructor. Used by the parent class'
        ``bind_expression`` method. """

    as_binary = "ST_AsEWKB"
    """ The "as binary" function to use. Used by the parent class'
        ``column_expression`` method. """

    ElementType = WKBElement
    """ The element class to use. Used by the parent class'
        ``result_processor`` method. """

    cache_ok = False
    """ Disable cache for this type. """


class Geography(_GISType):
    """The Geography type.

    Creating a geography column is done like this::

        Column(Geography(geometry_type='POINT', srid=4326))

    See :class:`geoalchemy2.types._GISType` for the list of arguments that can
    be passed to the constructor.
    """

    name = "geography"
    """ Type name used for defining geography columns in ``CREATE TABLE``. """

    from_text = "ST_GeogFromText"
    """ The ``FromText`` geography constructor. Used by the parent class'
        ``bind_expression`` method. """

    as_binary = "ST_AsBinary"
    """ The "as binary" function to use. Used by the parent class'
        ``column_expression`` method. """

    ElementType = WKBElement
    """ The element class to use. Used by the parent class'
        ``result_processor`` method. """

    cache_ok = False
    """ Disable cache for this type. """


class Raster(_GISType):
    """The Raster column type.

    Creating a raster column is done like this::

        Column(Raster)

    This class defines the ``result_processor`` method, so that raster values
    received from the database are converted to
    :class:`geoalchemy2.elements.RasterElement` objects.

    Args:
        spatial_index: Indicate if a spatial index should be created. Default is ``True``.
    """

    comparator_factory = BaseComparator
    """
    This is the way by which spatial operators and functions are
    defined for raster columns.
    """

    name = "raster"
    """ Type name used for defining raster columns in ``CREATE TABLE``. """

    from_text = "raster"
    """ The "from text" raster constructor. Used by the parent class'
        ``bind_expression`` method. """

    as_binary = "raster"
    """ The "as binary" function to use. Used by the parent class'
        ``column_expression`` method. """

    ElementType = RasterElement
    """ The element class to use. Used by the parent class'
        ``result_processor`` method. """

    cache_ok = False
    """ Disable cache for this type. """

    def __init__(self, spatial_index=True, from_text=None, name=None, nullable=True):
        # Enforce default values
        super(Raster, self).__init__(
            geometry_type=None,
            srid=-1,
            dimension=2,
            spatial_index=spatial_index,
            use_N_D_index=False,
            management=False,
            use_typmod=False,
            from_text=from_text,
            name=name,
            nullable=nullable,
        )
        self.extended = None

    @staticmethod
    def check_ctor_args(*args, **kwargs):
        return None, -1


class _DummyGeometry(Geometry):
    """A dummy type only used with SQLite."""

    def get_col_spec(self):
        return "GEOMETRY"


class CompositeType(UserDefinedType):
    """A composite type used by some spatial functions.

    A wrapper for :class:`geoalchemy2.elements.CompositeElement`, that can be
    used as the return type in PostgreSQL functions that return composite
    values.

    This is used as the base class of :class:`geoalchemy2.types.GeometryDump`.
    """

    typemap = {}
    """ Dictionary used for defining the content types and their
        corresponding keys. Set in subclasses. """

    class comparator_factory(UserDefinedType.Comparator):
        def __getattr__(self, key):
            try:
                type_ = self.type.typemap[key]
            except KeyError:
                raise AttributeError("Type '%s' doesn't have an attribute: '%s'" % (self.type, key))

            return CompositeElement(self.expr, key, type_)


class GeometryDump(CompositeType):
    """The return type for functions like ``ST_Dump``.

    The type consists in a path and a geom field.
    You should normally never use this class directly.
    """

    typemap = {"path": postgresql.ARRAY(Integer), "geom": Geometry}
    """ Dictionary defining the contents of a ``geometry_dump``. """

    cache_ok = True
    """ Enable cache for this type. """


# Register Geometry, Geography and Raster to SQLAlchemy's reflection subsystems.
postgresql_ischema_names["geometry"] = Geometry
postgresql_ischema_names["geography"] = Geography
postgresql_ischema_names["raster"] = Raster

sqlite_ischema_names["GEOMETRY"] = Geometry
sqlite_ischema_names["POINT"] = Geometry
sqlite_ischema_names["LINESTRING"] = Geometry
sqlite_ischema_names["POLYGON"] = Geometry
sqlite_ischema_names["MULTIPOINT"] = Geometry
sqlite_ischema_names["MULTILINESTRING"] = Geometry
sqlite_ischema_names["MULTIPOLYGON"] = Geometry
sqlite_ischema_names["CURVE"] = Geometry
sqlite_ischema_names["GEOMETRYCOLLECTION"] = Geometry
sqlite_ischema_names["RASTER"] = Raster


class SummaryStats(CompositeType):
    """Define the composite type returned by the function ST_SummaryStatsAgg."""

    typemap = {
        "count": Integer,
        "sum": Float,
        "mean": Float,
        "stddev": Float,
        "min": Float,
        "max": Float,
    }

    cache_ok = True
    """ Enable cache for this type. """
