#!/usr/bin/env python

"""
Database querying classes.

Copyright (C) 2006 Paul Boddie <paul@boddie.org.uk>

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License as published by the Free Software Foundation; either
version 2.1 of the License, or (at your option) any later version.

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA
"""

from sqltriples.Types import *

# Exceptions.

class NotSupportedError(Exception):
    pass

# Querying classes.

class AbstractQuery:

    table_columns = [("subject", "subject_type", 0), ("predicate", None, 0), ("object", "object_type", 0)]

    def _get_exposed_columns(self):
        return [(c, ct, 0) for (c, ct) in self.exposed_columns]

class Query(AbstractQuery):

    "An abstract class providing support for lazy query evaluation."

    def __init__(self, store, functions=None):
        self.store = store
        self.functions = functions or {}
        self.results = None
        self.result_type = self

    def _ensure(self):

        "Ensure the the query has been executed."

        if self.results is None:
            self.results = self._execute()

    def _execute(self):

        """
        Execute the query, generated by this object's 'get_query' method, and
        return the results in a form produced by this object's '_get_results'
        method.
        """

        cursor = self.store.connection.cursor()
        try:
            query, values, new_index = self.get_query(0)
            if self.store.debug:
                print self.store._pmarks(query), self.store._present(values)
            if values:
                cursor.execute(self.store._pmarks(query), self.store._present(values))
            else:
                cursor.execute(self.store._pmarks(query))
            results = cursor.fetchall()
        finally:
            cursor.close()
        return self._get_results(results)

    # Convenience methods for list-style access to results.

    def __getitem__(self, i):
        self._ensure()
        return self.results[i]

    def __len__(self):
        self._ensure()
        return len(self.results)

    def __repr__(self):
        self._ensure()
        return repr(self.results)

    def get_subquery(self, parent_columns, parent_index, index, match_first=0, match_parent=0, match_child=0):

        """
        Make a subquery for this object using the given 'parent_columns'
        (indicating the columns which connect with this object's exposed
        columns), a 'parent_index' (indicating the table index used in the
        parent query, within which the returned query clause shall be employed),
        and the current table 'index'.

        Return a query clause, associated values, and an updated table index (as
        a 3-tuple).
        """

        # Reserve a new table.

        child_index = index + 1

        # Support two styles of subquery.
        # parent_select_columns in (select child_select_columns ...)

        if self.store.adapter.supports_in_tuples:

            # Matching strategies:
            # Parent specifies columns: used mostly for matching predicates.

            if match_parent:
                child_select_columns = self._get_columns(parent_columns, child_index)
                parent_select_columns = self._get_columns(parent_columns, parent_index)

            # Child specifies columns: used in most cases to produce the expected
            # behaviour.

            elif match_child:
                child_columns = self._get_exposed_columns()
                child_select_columns = self._get_columns(child_columns, child_index)
                parent_select_columns = self._get_columns(child_columns, parent_index)

            # Match parent columns against available child columns: used for
            # chaining.

            elif match_first:
                child_columns = self.table_columns
                child_select_columns = self._get_columns(child_columns, child_index, limited_to=parent_columns)
                parent_select_columns = self._get_columns(parent_columns, parent_index)
            else:
                raise NotSupportedError, "match must be controlled by parent, child or first"

            operation = "(%s) in (%s)"
            operation_args = (parent_select_columns,)
            constraints = []

        # exists (select child_select_columns ... constraints)

        else:
            if match_parent:
                constraints = self._get_join_constraints(parent_columns, parent_columns, parent_index, child_index)
            elif match_child:
                child_columns = self._get_exposed_columns()
                constraints = self._get_join_constraints(child_columns, child_columns, parent_index, child_index)
            elif match_first:
                child_columns = self.table_columns
                constraints = self._get_join_constraints(parent_columns, child_columns, parent_index, child_index)
            else:
                raise NotSupportedError, "match must be controlled by parent, child or first"

            child_select_columns = "*"
            operation = "exists (%s)"
            operation_args = ()

        query = "select %s" % child_select_columns
        query += " " + self._get_table(child_index)
        joins, new_index = self._get_joins(child_index)
        query += " " + joins

        # Make the final query.

        query, values, new_index = self._get_query(query, child_index, new_index, " and ".join(constraints))
        return operation % (operation_args + (query,)), values, new_index

    def get_query(self, index, constraints=None):

        """
        Return a query clause, associated values, and an updated table index (as
        a 3-tuple) for this object's query defined using the given table 'index'
        and the supplied, optional 'constraints' text.
        """

        query = "select distinct %s" % self._get_columns(self._get_exposed_columns(), index)
        query += " " + self._get_table(index)
        joins, new_index = self._get_joins(index)
        query += " " + joins
        return self._get_query(query, index, new_index, constraints)

    def _apply_function(self, fn, column):
        if fn is not None:
            return fn.replace("_", column)
        else:
            return column

    def _get_table(self, index):
        return "from %s as %s" % (self.store._table(), self.store._table(index))

    def _get_columns(self, exposed_columns, index, limited_to=None):

        """
        Return a column selection for the 'exposed_columns' qualified by the
        given table 'index', optionally limited by the given 'limited_to'
        sequence of column definitions. The following style of text is produced:

        subject, subject_type, predicate, object, object_type, context
        """

        columns = []

        # Generate columns only for the exposed entries.

        if limited_to is not None:
            exposed_columns = exposed_columns[:len(limited_to)]

        for exposed_column, limiting_column in map(None, exposed_columns, limited_to or []):
            if exposed_column is None:
                break
            column, column_type, offset = exposed_column
            if limiting_column is not None:
                lcolumn, lcolumn_type, loffset = limiting_column
            columns.append(self._apply_function(self.functions.get(column), self.store._column(column, index + offset)))
            if column_type is not None and (limiting_column is None or lcolumn_type is not None):
                columns.append(self.store._column(column_type, index + offset))

        # Always add the context.

        columns.append(self.store._column("context", index))
        return ", ".join(columns)

    def _get_join_constraints(self, parent_columns, child_columns, parent_index, child_index):

        constraints = []

        # Generate columns only for the exposed entries.

        child_columns = child_columns[:len(parent_columns)]

        for child_column, parent_column in map(None, child_columns, parent_columns):
            if child_column is None:
                break
            column, column_type, offset = child_column
            if parent_column is not None:
                lcolumn, lcolumn_type, loffset = parent_column
                parent_str = self._apply_function(self.functions.get(lcolumn), self.store._column(lcolumn, parent_index + offset))
                child_str = self._apply_function(self.functions.get(column), self.store._column(column, child_index + loffset))
                constraints.append("%s = %s" % (parent_str, child_str))
                if column_type is not None and lcolumn_type is not None:
                    parent_type_str = self.store._column(lcolumn_type, parent_index + loffset)
                    child_type_str = self.store._column(column_type, child_index + offset)
                    constraints.append("%s = %s" % (parent_type_str, child_type_str))

        # Add the context.

        constraints.append("%s = %s" % (self.store._column("context", parent_index), self.store._column("context", child_index)))
        return constraints

class Results(Query):

    """
    A class representing a selection of results obtained by specifying a
    pattern.
    """

    def __init__(self, store, pattern, expression=None, functions=None, ordering=None, order_by=None, limit=None, partial=0):

        """
        Initialise the results object with the given 'store', 'pattern',
        optional 'expression', optional 'ordering' and 'order_by' descriptions
        and optional 'limit' criteria. Where a 'partial' result is requested,
        outer joins will be employed to find tuples which are shorter than the
        pattern.
        """

        self.pattern = pattern
        self.expression = expression
        self.ordering = ordering
        self.order_by = order_by
        self.limit = limit
        self.partial = partial
        Query.__init__(self, store, functions)

    def _get_query(self, select_clause, index, new_index, constraints=None, constraint_values=None):

        """
        Build the query conditions and return a usable query which starts with
        the supplied 'select_clause' and which uses the given table 'index'
        number as qualifier. The supplied 'new_index' specifies the latest
        table index in use, and the optional 'constraints' (a query string) and
        'constraint_values' (a list of values) provide additional qualification
        of the query.

        Return a 3-tuple containing the query text, query values and the updated
        table index.
        """

        # Add general conditions.

        if isinstance(self.pattern, Pattern) or isinstance(self.pattern, tuple):
            conditions, values, new_index = self.store._get_conditions(self.pattern, index, new_index)

        # Other things are not supported (previously we added a subquery for
        # each item in a sequence).

        else:
            raise NotSupportedError, self.pattern.__class__

        # For any supplied expressions, prepare a query which includes all the
        # subqueries.

        if self.expression is not None:
            subquery, subquery_values, new_index = self.expression.get_subquery(
                self._get_exposed_columns(), index, new_index, match_child=1)
            conditions.append(subquery)
            values += subquery_values

        # Add the constraints as an additional condition.

        if constraints:
            conditions.append(constraints)

        # Add all conditions to the query.

        if conditions:
            query = (select_clause + (" where %s" % " and ".join(conditions)))
        else:
            query = select_clause

        # Add the order clause.

        if self.ordering is not None:
            query += self._get_order_direction(index)

        # Add the limit clause.

        if self.limit:
            query += (" limit %s" % self.limit)

        return query, values + (constraint_values or []), new_index

    def _get_joins(self, index):

        """
        Add joins where the result needs them:
        ... inner join T1 on ... inner join T2 on ...
        ... left outer join T1 on ... left outer join T2 on ...
        """

        joins = []
        next_index = index
        if self.partial:
            join_type = "left outer"
        else:
            join_type = "inner"

        for i in range(0, len(self.pattern) - 3):
            this_index = index + i
            next_index = index + i + 1
            joins.append("%s join %s as %s on %s = %s and %s = %s" % (
                join_type,
                self.store._table(), self.store._table(next_index),
                self.store._column("object", this_index), self.store._column("subject", next_index),
                self.store._column("object_type", this_index), self.store._column("subject_type", next_index)
                ))

        return " ".join(joins), next_index

    def _get_order_direction(self, index):
        return " order by %s %s" % (self._get_order_column(index), self.ordering or "asc")

    def _get_order_extent(self):
        return len(self.pattern) - 3

    def _get_order_column(self, index):
        if self.order_by == "subject":
            return self.store._column("subject", 0) # NOTE: Verify this index!
        else:
            return self.store._column("object", index + self._get_order_extent())

class SingleResults(Results):

    "A class representing a selection of single value results."

    def _get_results(self, results):
        # Convert the fetched value using value and type information.
        return [self.store._instantiate(row[0], row[1]) for row in results]

class CountableResults(Results):

    "A class representing a count of some results."

    # Provided only for linkage/chaining.

    exposed_columns = [("subject", "subject_type"), ("predicate", None), ("object", "object_type")]

    def _get_results(self, results):
        return results

    def get_value(self):
        self._ensure()
        return self.results[0][0]

    def __getitem__(self, i):
        raise NotSupportedError, "__getitem__ for CountableResults is not supported"

    def __len__(self):
        # NOTE: Apparent Python limitation on the return value: must use int!
        return int(self.get_value())

    def __repr__(self):
        return str(self.get_value())

    def get_query(self, index):
        query = "select count(*)"
        query += " " + self._get_table(index)
        joins, new_index = self._get_joins(index)
        query += " " + joins
        return self._get_query(query, index, new_index)

    def get_subquery(self, parent_columns, parent_index, index, match_first=0, match_parent=0, match_child=0):
        raise NotSupportedError, "get_subquery for CountableResults is not supported"

class Subjects(SingleResults):

    "A class representing a selection of subjects."

    exposed_columns = [("subject", "subject_type")]

class Predicates(SingleResults):

    "A class representing a selection of predicates."

    exposed_columns = [("predicate", None)]

class Objects(SingleResults):

    "A class representing a selection of objects."

    exposed_columns = [("object", "object_type")]

class Triples(Results):

    "A class representing a selection of triples."

    exposed_columns = [("subject", "subject_type"), ("predicate", None), ("object", "object_type")]

    def _get_results(self, results):
        return [(self.store._instantiate(s, st), URIRef(p), self.store._instantiate(o, ot)) for (s, st, p, o, ot, c) in results]

class Tuples(Results):

    "A class representing a selection of triples."

    exposed_columns = [("subject", "subject_type"), ("predicate", None), ("object", "object_type")]

    def _get_results(self, results):

        "Return the 'results' for the retrieved tuples."

        final_results = []
        for row in results:
            # row[-1] is the context
            if not self.partial:
                final_results.append(
                    (self.store._instantiate(row[0], row[1]),) +
                    tuple([URIRef(p) for p in row[2:-3]]) +
                    (self.store._instantiate(row[-3], row[-2]),)
                    )
            else:
                # Convert...
                # s, p1, o1, p2, o2, p3, o3
                # ...into...
                # s, p1, p2, o2
                # ...when o3 is None.
                result = [self.store._instantiate(row[0], row[1])]
                last_object = None
                for i in range(2, len(row) - 1, 3):
                    p = row[i]
                    o, ot = row[i+1:i+3]
                    if o is None:
                        break
                    result.append(URIRef(p))
                    last_object = self.store._instantiate(o, ot)
                result.append(last_object)
                final_results.append(tuple(result))
        return final_results

    def _get_exposed_columns(self):

        "Return the extended exposed columns for the pattern."

        exposed = []
        if not self.partial:
            s, st = self.exposed_columns[0]
            exposed.append((s, st, 0))
            p, pt = self.exposed_columns[1]
            for i in range(0, len(self.pattern) - 2):
                exposed.append((p, pt, i))
            o, ot = self.exposed_columns[2]
            exposed.append((o, ot, len(self.pattern) - 3))
        else:
            s, st = self.exposed_columns[0]
            exposed.append((s, st, 0))
            p, pt = self.exposed_columns[1]
            o, ot = self.exposed_columns[2]
            for i in range(0, len(self.pattern) - 2):
                exposed.append((p, pt, i))
                exposed.append((o, ot, i))
        return exposed

# Special combining query classes.

class Negation:

    "A class representing the negation of other results."

    def get_subquery(self, parent_columns, parent_index, index, match_first=0, match_parent=0, match_child=0):
        query, values, new_index = self.result.get_subquery(
            parent_columns, parent_index, index, match_first, match_parent, match_child)
        return "not %s" % query, values, new_index

    def get_query(self, index):
        query = "select distinct %s" % self._get_columns(self._get_exposed_columns(), index)
        query += " " + self._get_table(index)
        joins, new_index = self._get_joins(index)
        query += " " + joins
        parent_columns = self._get_exposed_columns()
        subquery, subquery_values, new_index = self.get_subquery(parent_columns, index, new_index, match_child=1)
        return self._get_query(query, index, new_index, constraints=subquery, constraint_values=subquery_values)

class NegationOfSubjects(Negation, Subjects):
    def __init__(self, store, result):
        Subjects.__init__(self, store, ALL)
        self.result = result

class NegationOfPredicates(Negation, Predicates):
    def __init__(self, store, result):
        Predicates.__init__(self, store, ALL)
        self.result = result

class NegationOfObjects(Negation, Objects):
    def __init__(self, store, result):
        Objects.__init__(self, store, ALL)
        self.result = result

class NegationOfTriples(Negation, Triples):
    def __init__(self, store, result):
        Triples.__init__(self, store, ALL)
        self.result = result

class NegationOfTuples(Negation, Tuples):
    def __init__(self, store, result, pattern=ALL):
        Tuples.__init__(self, store, pattern)
        self.result = result

class CombinedResults:

    def get_subquery(self, parent_columns, parent_index, index, match_first=0, match_parent=0, match_child=0):
        new_index = index
        subqueries = []
        values = []
        for query in self.queries:
            subquery, subquery_values, new_index = query.get_subquery(
                parent_columns, parent_index, new_index, match_first, match_parent, match_child)
            subqueries.append(subquery)
            values += subquery_values
        final_query = (" %s " % self.operator).join(subqueries)
        if self.in_parentheses:
            final_query = "(" + final_query + ")"
        return final_query, values, new_index

    def get_query(self, index):
        query = "select distinct %s" % self._get_columns(self._get_exposed_columns(), index)
        query += " " + self._get_table(index)
        joins, new_index = self._get_joins(index)
        query += " " + joins
        parent_columns = self._get_exposed_columns()
        subquery, subquery_values, new_index = self.get_subquery(parent_columns, index, new_index, match_child=1)
        return self._get_query(query, index, new_index, constraints=subquery, constraint_values=subquery_values)

class Conjunction(CombinedResults):
    operator = "and"
    in_parentheses = 1

class Disjunction(CombinedResults):
    operator = "or"
    in_parentheses = 1

class ConjunctionOfSubjects(Conjunction, Subjects):
    def __init__(self, store, queries):
        Subjects.__init__(self, store, ALL)
        self.queries = queries

class ConjunctionOfPredicates(Conjunction, Predicates):
    def __init__(self, store, queries):
        Predicates.__init__(self, store, ALL)
        self.queries = queries

class ConjunctionOfObjects(Conjunction, Objects):
    def __init__(self, store, queries):
        Objects.__init__(self, store, ALL)
        self.queries = queries

class ConjunctionOfTriples(Conjunction, Triples):
    def __init__(self, store, queries):
        Triples.__init__(self, store, ALL)
        self.queries = queries

class ConjunctionOfTuples(Conjunction, Tuples):
    def __init__(self, store, queries, pattern=ALL):
        Tuples.__init__(self, store, pattern)
        self.queries = queries

class DisjunctionOfSubjects(Disjunction, Subjects):
    def __init__(self, store, queries):
        Subjects.__init__(self, store, ALL)
        self.queries = queries

class DisjunctionOfPredicates(Disjunction, Predicates):
    def __init__(self, store, queries):
        Predicates.__init__(self, store, ALL)
        self.queries = queries

class DisjunctionOfObjects(Disjunction, Objects):
    def __init__(self, store, queries):
        Objects.__init__(self, store, ALL)
        self.queries = queries

class DisjunctionOfTriples(Disjunction, Triples):
    def __init__(self, store, queries):
        Triples.__init__(self, store, ALL)
        self.queries = queries

class DisjunctionOfTuples(Disjunction, Tuples):
    def __init__(self, store, queries, pattern=ALL):
        Tuples.__init__(self, store, pattern)
        self.queries = queries

# Whole query classes.

class CombinedTopLevelResults(CombinedResults):
    def get_subquery(self, parent_columns, parent_index, index, match_first=0, match_parent=0, match_child=0):
        raise NotSupportedError, "get_subquery"

    def get_query(self, index):
        subqueries = []
        values = []
        for query in self.queries:
            subquery, subquery_values, new_index = query.get_query(index)
            subqueries.append(subquery)
            values += subquery_values
        final_query = (" %s " % self.operator).join(subqueries)
        if self.in_parentheses:
            final_query = "(" + final_query + ")"
        return final_query, values, new_index

class Intersection(CombinedTopLevelResults):
    operator = "intersect"
    in_parentheses = 0

class IntersectionOfSubjects(Intersection, Subjects):
    def __init__(self, store, queries):
        Subjects.__init__(self, store, ALL)
        self.queries = queries

class IntersectionOfPredicates(Intersection, Predicates):
    def __init__(self, store, queries):
        Predicates.__init__(self, store, ALL)
        self.queries = queries

class IntersectionOfObjects(Intersection, Objects):
    def __init__(self, store, queries):
        Objects.__init__(self, store, ALL)
        self.queries = queries

class IntersectionOfTriples(Intersection, Triples):
    def __init__(self, store, queries):
        Triples.__init__(self, store, ALL)
        self.queries = queries

class IntersectionOfTuples(Intersection, Tuples):
    def __init__(self, store, queries, pattern=ALL):
        Tuples.__init__(self, store, pattern)
        self.queries = queries

# Convenience classes.

class Subject(AbstractQuery):

    "A class representing a more conveniently accessible subject."

    exposed_columns = [("subject", "subject_type")]

    def __init__(self, store, subject):

        """
        Initialise the subject instance with the given 'store' and 'subject'
        identifier.
        """

        self.store = store
        self.subject = subject

    def __getitem__(self, predicate):
        return self.store.objects(self.subject, predicate)

    def __delitem__(self, predicate):
        if isinstance(predicate, tuple):
            predicate, object = predicate
            self.store.remove((self.subject, predicate, object))
        else:
            self.store.remove((self.subject, predicate, None))

    def __setitem__(self, predicate, object):
        self.store.add((self.subject, predicate, object))

    def keys(self):
        return self.store.predicates(self.subject)

    def values(self):
        return self.store.triples((self.subject, None, None))

    # Querying methods.

    def get_child_query(self, index):
        return self.store.subjects(pattern=(self.subject, None, None)).get_child_query(index)

    def get_query(self, index, constraints=None):
        return self.store.subjects(pattern=(self.subject, None, None)).get_query(index, constraints)

    def get_subquery(self, parent_columns, parent_index, index, match_first=0, match_parent=0, match_child=0):
        return self.store.subjects(pattern=(self.subject, None, None)).get_subquery(
            parent_columns, parent_index, index, match_first, match_parent, match_child)

# vim: tabstop=4 expandtab shiftwidth=4
