# Copyright (C) 2003, 2004, 2005 by Intevation GmbH
# Authors:
# Martin Mueller <mmueller@intevation.de>
# Bernhard Herzog <bh@intevation.de>
#
# This program is free software under the GPL (>=v2)
# Read the file COPYING coming with the software for details.

"""Basic interface to a PostGIS database"""

from __future__ import generators

try:
    import psycopg2
    from psycopg2 import extensions
except ImportError:
    psycopg2 = None

import table
import wellknowntext

from data import SHAPETYPE_POLYGON, SHAPETYPE_ARC, SHAPETYPE_POINT, RAW_WKT

def has_postgis_support():
    """Return whether this Thuban instance supports PostGIS connections

    Having PostGIS support means that the psycopg module can be
    imported.
    """
    return psycopg2 is not None

def psycopg_version():
    return psycopg2.__version__

if psycopg2 is not None:
    # type_map maps psycopg type objects.  It's a list of pairs since
    # the psycopg type objects are unhashable.
    type_map = [(psycopg2.STRING, table.FIELDTYPE_STRING),
                (extensions.INTEGER, table.FIELDTYPE_INT),
                (psycopg2.ROWID, table.FIELDTYPE_INT),
                (extensions.FLOAT, table.FIELDTYPE_DOUBLE)]

    # _raw_type_map maps the postgresql type constants to Thuban type
    # constants.  This is very low level and postgresql specific and
    # should be used only when necessary.
    _raw_type_map = {}
    def _fill_raw_type_map():
        for psycopg_type, thuban_type in type_map:
            for value in psycopg_type.values:
                _raw_type_map[value] = thuban_type
    _fill_raw_type_map()


def quote_identifier(ident):
    """Return a quoted version of the identifier ident.

    The return value is a string that can be put directly into an SQL
    statement.  The quoted identifier is surrounded by double quotes and
    any double quotes already in the input value are converted to two
    double quotes.  Examples:

    >>> quote_identifier("abc\"def")
    '"abc""def"'
    >>> quote_identifier("abc def")
    '"abc def"'
    """
    return '"' + '""'.join(ident.split('"')) + '"'


class ConnectionError(Exception):

    """Class for exceptions occurring when establishing a Databse connection"""


class PostGISConnection:

    """Represent a PostGIS database

    A PostGISConnection instance has the following public attributes:

    dbname -- The name of the database
    host, port -- Host and port to connect to
    user -- The user name to connect as.

    All of these attributes are strings and may be empty strings to
    indicate default values.
    """

    def __init__(self, dbname, host="", user="", password="", dbtype="",
                 port=""):
        self.dbname = dbname
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.dbtype = dbtype
        self.connect()

    def connect(self):
        """Internal: Establish the database connection"""
        params = []
        for name in ("host", "port", "dbname", "user", "password"):
            val = getattr(self, name)
            if val:
                params.append("%s=%s" % (name, val))
        try:
            self.connection = psycopg2.connect(" ".join(params))
        except psycopg2.OperationalError, val:
            raise ConnectionError(str(val))

        # Use autocommit mode.  For simple reading of the database it's
        # sufficient and we don't have to care much about error
        # handling.  Without autocommit, an errors during a cursor's
        # execute method requires a rollback on the connection,
        # otherwise later queries with the same or other cursors sharing
        # the same connection will lead to further errors ("ERROR:
        # current transaction is aborted, commands ignored until end of
        # transaction block")
        self.connection.set_isolation_level(0)

        # determine the OID for the geometry type. This is PostGIS
        # specific.
        cursor = self.connection.cursor()
        cursor.execute("SELECT OID, typname FROM pg_type WHERE"
                       +" typname = 'geometry'")
        row = cursor.fetchone()
        self.connection.commit()
        if row is not None:
            self.geometry_type = row[0]
        else:
            raise ValueError("Can't determine postgres type of geometries")

    def BriefDescription(self):
        """Return a brief, one-line description of the connection

        The return value is suitable for a list box of all database
        connections.
        """
        return ("postgis://%(user)s@%(host)s:%(port)s/%(dbname)s"
                % self.__dict__)

    def MatchesParameters(self, parameters):
        """Return whether the connection matches the dictionary of parameters

        Return whether instatiating the connection with the given
        parameters would establish essentially the same connection as
        self. The connection is essentially the same if the same
        database (identified by host, port and databasename) is accessed
        as the same user.
        """
        return (parameters["host"] == self.host
                and parameters["port"] == self.port
                and parameters["dbname"] == self.dbname
                and parameters["user"] == self.user)

    def Close(self):
        """Close the database connection"""
        self.connection.close()

    def GeometryTables(self):
        """Return a list with the names of all tables with a geometry column"""

        # The query is basically taken from the psql v. 7.2.1.  When
        # started with -E it prints the queries used for internal
        # commands such as \d, which does mostly what we need here.
        cursor = self.connection.cursor()
        cursor.execute("SELECT c.relname FROM pg_class c"
                       " WHERE c.relkind IN ('r', 'v')"
                             # Omit the system tables
                             " AND c.relname !~ '^pg_'"
                             # Omit the special PostGIS tables
                             " AND c.relname NOT IN ('geometry_columns',"
                                                   " 'spatial_ref_sys')"
                            " AND %s in (SELECT a.atttypid FROM pg_attribute a"
                                       " WHERE a.attrelid = c.oid)"
                       " ORDER BY c.relname;", (self.geometry_type,))
        result = [row[0] for row in cursor.fetchall()]
        self.connection.commit()
        return result

    def table_columns(self, tablename):
        """Experimental: return information about the columns of a table

        Return value is a list of (name, type) pairs where name is the
        name of the column and type either one of the field type columns
        or the string 'geometry' indicating a geometry column.

        The intended use of this method is for table selection dialogs
        which need to determine which columns are usable as id or
        geometry columns respectively.  Suitable id columns will have
        type FIELDTYPE_INT and geometry columns will have 'geometry'.
        """
        result = []
        cursor = self.connection.cursor()

        # This query is taken basically from the \d command of psql
        # 7.2.1
        cursor.execute("SELECT a.attname, a.atttypid, a.attnum"
                       " FROM pg_class c, pg_attribute a"
                            " WHERE c.relname = %s AND a.attrelid = c.oid"
                       " ORDER BY a.attnum;", (tablename,))

        for row in cursor.fetchall():
            col_name, col_type, col_attnum = row
            col = None
            if col_attnum < 1:
                # It's a system column.  Only the OID is interesting
                # here
                if col_name == "oid":
                    col = (col_name, _raw_type_map[col_type])
            else:
                # If it's an integer
                thuban_type = _raw_type_map.get(col_type)
                if thuban_type is not None:
                    col = (col_name, thuban_type)
                elif row[1] == self.geometry_type:
                    col = (col_name, "geometry")
            if col is not None:
                result.append(col)

        return result

    def cursor(self):
        """Return a DB API 2.0 cursor for the database"""
        return self.connection.cursor()



class PostGISColumn:

    """Column description for a PostGISTable

    In addition to the normal column object attributes name, type and
    index, PostGISColumn objects have a quoted_name attribute which
    contains a quoted version of name for use in SQL statements. The
    quoted_name attribute is mainly intended for internal use by the
    PostGISTable class.
    """

    def __init__(self, name, type, index):
        self.name = name
        self.quoted_name = quote_identifier(name)
        self.type = type
        self.index = index


class PostGISTable:

    """A Table in a PostGIS database

    A PostgreSQL table may contain columns with types not (yet)
    supported by Thuban. Instances of this class ignore those columns
    and pretend they don't exist, i.e. they won't show up in the column
    descriptions returned by Columns() and other methods.
    """

    def __init__(self, db, tablename, id_column = None):
        """Initialize the PostGISTable.

        The db parameter should be an instance of PostGISConnection and
        tablename the name of a table in the database represented by db.

        The id_column parameter should be the name of a column in the
        table that can be used to identify rows.  The column must have
        the type integer and be unique and not null.

        For backwards compatibility reasons, the id_column parameter is
        optional.  If not given the table must have a column called
        'gid' which is used as the id_column.  New code should always
        provide this parameter.
        """
        self.db = db
        self.tablename = tablename
        # Tablename quoted for use in SQL statements.
        self.quoted_tablename = quote_identifier(tablename)

        if not id_column:
            id_column = "gid"
        self.id_column = id_column
        # id column name quoted for use in SQL statements.
        self.quoted_id_column = quote_identifier(id_column)

        # Map column names and indices to column objects.
        self.column_map = {}

        self._fetch_table_information()

    def _fetch_table_information(self):
        """Internal: Update information about the table"""
        self.columns = []
        cursor = self.db.cursor()
        cursor.execute("SELECT * FROM %s WHERE 0=1" % self.quoted_tablename)
        description = cursor.description

        for i in range(len(description)):
            col = self._create_col_from_description(i, description[i])
            if col is not None:
                self.columns.append(col)

        for col in self.columns:
            self.column_map[col.name] = col
            self.column_map[col.index] = col

        # Build query string for ReadRowAsDict
        self.query_stmt = ("SELECT %s from %s"
                           % (", ".join([col.quoted_name
                                         for col in self.columns]),
                              self.quoted_tablename))

    def _create_col_from_description(self, index, description):
        """Return the column object for the column described by description

        The parameter index is the index of the column.  The description
        is a sequence taken from the cursor's description attribute for
        the column.  That means description[0] is the name of the column
        and description[1] the type.

        Return None if the column can't be represented for some reason,
        e.g. because its type is not yet supported or needs to be
        treated in some special way.  Derived classes may extend this
        method.
        """
        for pgtyp, tabletyp in type_map:
            if pgtyp == description[1]:
                return PostGISColumn(description[0], tabletyp,
                                     len(self.columns))
        return None

    def DBConnection(self):
        """Return the dbconnection used by the table"""
        return self.db

    def IDColumn(self):
        """Return the column description object for the id column.

        If the oid column was used as the id column, the return value is
        not one of the regular column objects that would be returned by
        e.g. the Column() method, but it still has meaningful name
        attribute.
        """
        if self.id_column == "oid":
            return PostGISColumn(self.id_column, table.FIELDTYPE_INT, None)
        return self.column_map[self.id_column]

    def TableName(self):
        """Return the name of the table in the database"""
        return self.tablename

    def Title(self):
        """Return the title of the table.

        The title is currently fixed and equal to the tablename
        """
        return self.tablename

    def Dependencies(self):
        """Return an empty tuple because a PostGISTable depends on nothing else
        """
        return ()

    def NumColumns(self):
        return len(self.columns)

    def Columns(self):
        return self.columns

    def Column(self, col):
        return self.column_map[col]

    def HasColumn(self, col):
        return self.column_map.has_key(col)

    def NumRows(self):
        cursor = self.db.cursor()
        cursor.execute("SELECT count(*) FROM %s" % self.quoted_tablename)
        return cursor.fetchone()[0]

    def RowIdToOrdinal(self, gid):
        """Return the row ordinal given its id"""
        cursor = self.db.cursor()
        cursor.execute("SELECT count(*) FROM %s WHERE %s < %s;"
                       % (self.quoted_tablename, self.quoted_id_column, gid))
        return cursor.fetchone()[0]

    def RowOrdinalToId(self, num):
        """Return the rowid for given its ordinal"""
        cursor = self.db.cursor()
        cursor.execute("SELECT %s FROM %s LIMIT 1 OFFSET %s;"
                       % (self.quoted_id_column, self.quoted_tablename, num))
        return cursor.fetchone()[0]

    def ReadRowAsDict(self, row, row_is_ordinal = 0):
        cursor = self.db.cursor()
        if row_is_ordinal:
            stmt = self.query_stmt + " LIMIT 1 OFFSET %s" % row
        else:
            stmt = self.query_stmt + " WHERE %s = %s" % (self.quoted_id_column,
                                                         row)
        cursor.execute(stmt)
        result = {}
        for col, value in zip(self.columns, cursor.fetchone()):
            result[col.name] = value
        return result

    def ReadValue(self, row, col, row_is_ordinal = 0):
        cursor = self.db.cursor()
        if row_is_ordinal:
            stmt = ("SELECT %s FROM %s LIMIT 1 OFFSET %s" %
                    (self.column_map[col].quoted_name, self.quoted_tablename,
                     row))
        else:
            stmt = ("SELECT %s FROM %s WHERE %s = %s" %
                    (self.column_map[col].quoted_name, self.quoted_tablename,
                     self.quoted_id_column, row))
        cursor.execute(stmt)
        return cursor.fetchone()[0]

    def ValueRange(self, col):
        cursor = self.db.cursor()
        name = self.column_map[col].quoted_name
        cursor.execute("SELECT min(%s), max(%s) FROM %s" %
                       (name, name, self.quoted_tablename))
        return tuple(cursor.fetchone())

    def UniqueValues(self, col):
        cursor = self.db.cursor()
        name = self.column_map[col].quoted_name
        cursor.execute("SELECT %s FROM %s GROUP BY %s" %
                       (name, self.quoted_tablename, name))
        return [row[0] for row in cursor.fetchall()]

    def SimpleQuery(self, left, comparison, right):
        if comparison not in ("==", "!=", "<", "<=", ">=", ">"):
            raise ValueError("Comparison operator %r not allowed" % comparison)

        if comparison == "==":
            comparison = "="

        if isinstance(right, PostGISColumn):
            right_template = right.quoted_name
            params = ()
        else:
            right_template = "%s"
            params = (right,)

        query = "SELECT %s FROM %s WHERE %s %s %s ORDER BY %s;" \
                % (self.quoted_id_column, self.quoted_tablename,
                   left.quoted_name, comparison, right_template,
                   self.quoted_id_column)

        cursor = self.db.cursor()
        cursor.execute(query, params)
        result = []
        while 1:
            row = cursor.fetchone()
            if row is None:
                break
            result.append(row[0])
        return result


class PostGISShape:

    def __init__(self, shapeid, data):
        self.shapeid = shapeid
        self.data = data

    def compute_bbox(self):
        """
        Return the bounding box of the shape as a tuple (minx,miny,maxx,maxy)
        """
        xs = []
        ys = []
        for part in self.Points():
            for x, y in part:
                xs.append(x)
                ys.append(y)
        return (min(xs), min(ys), max(xs), max(ys))

    def ShapeID(self):
        return self.shapeid

    def Points(self):
        return wellknowntext.parse_wkt_thuban(self.data)

    def RawData(self):
        return self.data


shapetype_map = {"POLYGON": SHAPETYPE_POLYGON,
                 "MULTIPOLYGON": SHAPETYPE_POLYGON,
                 "LINESTRING": SHAPETYPE_ARC,
                 "MULTILINESTRING": SHAPETYPE_ARC,
                 "POINT": SHAPETYPE_POINT}


class PostGISShapeStore(PostGISTable):

    """Shapestore interface to a table in a PostGIS database"""

    def __init__(self, db, tablename, id_column = "gid",
                 geometry_column = None):
        """Initialize the PostGISShapeStore.

        The db parameter should be an instance of PostGISConnection and
        tablename the name of a table in the database represented by db.

        The id_column parameter should be the name of a column in the
        table that can be used to identify rows.  The column must have
        the type integer and be unique and not null.

        The geometry_column paramter, if given, should be the name of
        the geometry column to use.  If the name given is not a geometry
        column, raise a ValueError.

        If no geometry_column is given, the table must have exactly one
        geometry column.  If it has more than one and the
        geometry_column is not given, a ValueError will be raised.
        """
        self.geometry_column = geometry_column
        self.geometry_column_was_given = geometry_column is not None
        PostGISTable.__init__(self, db, tablename, id_column)

        # For convenience, we have a quoted version of the geometry
        # column in self.quoted_geo_col
        self.quoted_geo_col = quote_identifier(self.geometry_column)

    def _fetch_table_information(self):
        """Extend inherited method to retrieve the SRID and shape type"""
        PostGISTable._fetch_table_information(self)

        # First, try to get it from the geometry_columns table.
        cursor = self.db.cursor()
        cursor.execute("SELECT srid, type FROM geometry_columns"
                       " WHERE f_table_name = %s AND f_geometry_column = %s",
                       (self.tablename, self.geometry_column))
        row = cursor.fetchone()
        if row is not None:
            self.srid = row[0]
            self.shape_type = shapetype_map.get(row[1])
            return

        # The table is probably really a view and thus not in
        # geometry_columns.  Use a different approach
        cursor = self.db.cursor()
        cursor.execute("SELECT DISTINCT SRID(%s) FROM %s;" %
                       (quote_identifier(self.geometry_column),
                        quote_identifier(self.tablename)))
        row = cursor.fetchone()
        if row is not None:
            self.srid = row[0]
            # Try to see whether there's another one
            row = cursor.fetchone()
            if row is not None:
                # There are at least two different srids.  We don't
                # support that
                self.srid = None

        cursor = self.db.cursor()
        cursor.execute("SELECT DISTINCT GeometryType(%s) FROM %s;"
                       % (quote_identifier(self.geometry_column),
                          quote_identifier(self.tablename)))
        row = cursor.fetchone()
        if row is not None:
            self.shape_type = shapetype_map.get(row[0])
            # Try to see whether there's another one
            row = cursor.fetchone()
            if row is not None:
                # There are at least two different srids.  We don't
                # support that
                self.shape_type = None

    def _create_col_from_description(self, index, description):
        """Extend the inherited method to find geometry columns

        If the column indicated by the parameters is a geometry column,
        record its name in self.geometry_column and a quoted version in
        self.quoted_geo_col.  In any case return the return value of the
        inherited method.
        """
        col = PostGISTable._create_col_from_description(self, index,
                                                        description)
        col_name, col_type = description[:2]
        if self.geometry_column_was_given:
            if (col_name == self.geometry_column
                and col_type != self.db.geometry_type):
                raise TypeError("Column %s in %s is not a geometry column"
                                % (self.geometry_column, self.tablename))
        else:
            if col is None:
                if description[1] == self.db.geometry_type:
                    # The column is a geometry column.  If the name of
                    # the geometry column was not given to the
                    # constructor, and we encounter two geometry
                    # columns, raise a value error
                    if self.geometry_column is None:
                        self.geometry_column = description[0]
                    else:
                        raise TypeError("Table %s has two geometry columns"
                                        " and no column name was given"
                                        % (self.tablename,))
        return col

    def Table(self):
        """Return self since a PostGISShapeStore is its own table."""
        return self

    def OrigShapeStore(self):
        """Return None since the PostGISShapeStore is not derived from another
        """
        return None

    def GeometryColumn(self):
        """Return the column description object for the geometry column

        There's currently no FIELDTYPE constant for this column, so the
        return value is not a regular column object that could also be
        returned from e.g. the Column() method.  Only the name attribute
        of the return value is meaningful at the moment.
        """
        return PostGISColumn(self.geometry_column, None, None)

    def ShapeType(self):
        """Return the type of the shapes in the shapestore."""
        return self.shape_type

    def RawShapeFormat(self):
        """Return the raw data format of the shape data.

        For the PostGISShapeStore this is RAW_WKT.
        """
        return RAW_WKT

    def NumShapes(self):
        # The number of shapes is the same as the number of rows,
        # assuming that the geometry can't be NULL.
        return self.NumRows()

    def BoundingBox(self):
        """Return the bounding box of all shapes in the postgis table"""
        minx = miny = maxx = maxy = None
        x=[]
        y=[]
        cursor = self.db.cursor()
        try:
            # Using the extent function is postgis specific. An OGC
            # Simple Features compliant solution would be to use a query
            # like "SELECT AsText(Envelope(the_geom)) FROM mytable;" and
            # calculate the bounding box by hand from that
            cursor.execute("SELECT extent(%s) FROM %s;"
                           % (self.quoted_geo_col, self.quoted_tablename))
            result = cursor.fetchone()
            if result[0]:
                (minx, miny), (maxx, maxy) \
                      = wellknowntext.parse_wkt_thuban(result[0])[0]
                return (minx, miny, maxx, maxy)
        finally:
            cursor.close()

    def Shape(self, shapeid):
        cursor = self.db.cursor()
        cursor.execute("SELECT AsText(%s) FROM %s WHERE %s = %s"
                       % (self.quoted_geo_col, self.quoted_tablename,
                          self.quoted_id_column, shapeid))
        wkt = cursor.fetchone()[0]
        cursor.close()
        return PostGISShape(shapeid, wkt)

    def AllShapes(self):
        cursor = self.db.cursor()
        cursor.execute("SELECT %s, AsText(%s) FROM %s ORDER BY %s"
                       % (self.quoted_id_column, self.quoted_geo_col,
                          self.quoted_tablename, self.quoted_id_column))
        while 1:
            result = cursor.fetchone()
            if result is None:
                return
            yield PostGISShape(result[0], result[1])


    def ShapesInRegion(self, bbox):
        """Generate all shapes overlapping the region given by bbox."""
        # IMPORTANT:This will work for PostGIS < 0.8
        left, bottom, right, top = bbox
        geom = (("POLYGON((" + ", ".join(["%f %f"] * 5) + "))")
                % (left, bottom, left, top, right, top, right, bottom,
                   left, bottom))
        cursor = self.db.cursor()
        cursor.execute("SELECT %(gid)s, AsText(%(geom)s) FROM %(table)s"
                     " WHERE %(geom)s && GeometryFromText('%(box)s', %(srid)d)"
                       " ORDER BY %(gid)s"
                       % {"table": self.quoted_tablename,
                          "geom": self.quoted_geo_col,
                          "gid": self.quoted_id_column,
                          "box": geom,
                          "srid": self.srid})
        while 1:
            result = cursor.fetchone()
            if result is None:
                return
            yield PostGISShape(result[0], result[1])
