# copyright 2003-2013 LOGILAB S.A. (Paris, FRANCE), all rights reserved.

# contact http://www.logilab.fr/ -- mailto:contact@logilab.fr

#

# This file is part of CubicWeb.

#

# CubicWeb 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.

#

# CubicWeb 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 CubicWeb.  If not, see <http://www.gnu.org/licenses/>.

"""RQL rewriting utilities : insert rql expression snippets into rql syntax

tree.

This is used for instance for read security checking in the repository.

"""

__docformat__ = "restructuredtext en"

from rql import nodes as n, stmts, TypeResolverException

from rql.utils import common_parent

from yams import BadSchemaDefinition

from logilab.common import tempattr

from logilab.common.graph import has_path

from cubicweb import Unauthorized

def cleanup_solutions(rqlst, solutions):

    for sol in solutions:

        for vname in list(sol):

            if not (vname in rqlst.defined_vars or vname in rqlst.aliases):

                del sol[vname]

def add_types_restriction(schema, rqlst, newroot=None, solutions=None):

    if newroot is None:

        assert solutions is None

        if hasattr(rqlst, '_types_restr_added'):

            return

        solutions = rqlst.solutions

        newroot = rqlst

        rqlst._types_restr_added = True

    else:

        assert solutions is not None

        rqlst = rqlst.stmt

    eschema = schema.eschema

    allpossibletypes = {}

    for solution in solutions:

        for varname, etype in solution.iteritems():

            # XXX not considering aliases by design, right ?

            if varname not in newroot.defined_vars or eschema(etype).final:

                continue

            allpossibletypes.setdefault(varname, set()).add(etype)

    # XXX could be factorized with add_etypes_restriction from rql 0.31

    for varname in sorted(allpossibletypes):

        var = newroot.defined_vars[varname]

        stinfo = var.stinfo

        if stinfo.get('uidrel') is not None:

            continue # eid specified, no need for additional type specification

        try:

            typerel = rqlst.defined_vars[varname].stinfo.get('typerel')

        except KeyError:

            assert varname in rqlst.aliases

            continue

        if newroot is rqlst and typerel is not None:

            mytyperel = typerel

        else:

            for vref in var.references():

                rel = vref.relation()

                if rel and rel.is_types_restriction():

                    mytyperel = rel

                    break

            else:

                mytyperel = None

        possibletypes = allpossibletypes[varname]

        if mytyperel is not None:

            if mytyperel.r_type == 'is_instance_of':

                # turn is_instance_of relation into a is relation since we've

                # all possible solutions and don't want to bother with

                # potential is_instance_of incompatibility

                mytyperel.r_type = 'is'

                if len(possibletypes) > 1:

                    node = n.Function('IN')

                    for etype in possibletypes:

                        node.append(n.Constant(etype, 'etype'))

                else:

                    node = n.Constant(etype, 'etype')

                comp = mytyperel.children[1]

                comp.replace(comp.children[0], node)

            else:

                # variable has already some strict types restriction. new

                # possible types can only be a subset of existing ones, so only

                # remove no more possible types

                for cst in mytyperel.get_nodes(n.Constant):

                    if not cst.value in possibletypes:

                        cst.parent.remove(cst)

        else:

            # we have to add types restriction

            if stinfo.get('scope') is not None:

                rel = var.scope.add_type_restriction(var, possibletypes)

            else:

                # tree is not annotated yet, no scope set so add the restriction

                # to the root

                rel = newroot.add_type_restriction(var, possibletypes)

            stinfo['typerel'] = rel

        stinfo['possibletypes'] = possibletypes

def remove_solutions(origsolutions, solutions, defined):

    """when a rqlst has been generated from another by introducing security

    assertions, this method returns solutions which are contained in orig

    solutions

    """

    newsolutions = []

    for origsol in origsolutions:

        for newsol in solutions[:]:

            for var, etype in origsol.items():

                try:

                    if newsol[var] != etype:

                        try:

                            defined[var].stinfo['possibletypes'].remove(newsol[var])

                        except KeyError:

                            pass

                        break

                except KeyError:

                    # variable has been rewritten

                    continue

            else:

                newsolutions.append(newsol)

                solutions.remove(newsol)

    return newsolutions

def _add_noinvariant(noinvariant, restricted, select, nbtrees):

    # a variable can actually be invariant if it has not been restricted for

    # security reason or if security assertion hasn't modified the possible

    # solutions for the query

    for vname in restricted:

        try:

            var = select.defined_vars[vname]

        except KeyError:

            # this is an alias

            continue

        if nbtrees != 1 or len(var.stinfo['possibletypes']) != 1:

            noinvariant.add(var)

def _expand_selection(terms, selected, aliases, select, newselect):

    for term in terms:

        for vref in term.iget_nodes(n.VariableRef):

            if not vref.name in selected:

                select.append_selected(vref)

                colalias = newselect.get_variable(vref.name, len(aliases))

                aliases.append(n.VariableRef(colalias))

                selected.add(vref.name)

def _has_multiple_cardinality(etypes, rdef, ttypes_func, cardindex):

    """return True if relation definitions from entity types (`etypes`) to

    target types returned by the `ttypes_func` function all have single (1 or ?)

    cardinality.

    """

    for etype in etypes:

        for ttype in ttypes_func(etype):

            if rdef(etype, ttype).cardinality[cardindex] in '+*':

                return True

    return False

def iter_relations(stinfo):

    # this is a function so that test may return relation in a predictable order

    return stinfo['relations'] - stinfo['rhsrelations']

class Unsupported(Exception):

    """raised when an rql expression can't be inserted in some rql query

    because it create an unresolvable query (eg no solutions found)

    """

class RQLRewriter(object):

    """insert some rql snippets into another rql syntax tree

    this class *isn't thread safe*

    """

    def __init__(self, session):

        self.session = session

        vreg = session.vreg

        self.schema = vreg.schema

        self.annotate = vreg.rqlhelper.annotate

        self._compute_solutions = vreg.solutions

    def compute_solutions(self):

        self.annotate(self.select)

        try:

            self._compute_solutions(self.session, self.select, self.kwargs)

        except TypeResolverException:

            raise Unsupported(str(self.select))

        if len(self.select.solutions) < len(self.solutions):

            raise Unsupported()

    def insert_local_checks(self, select, kwargs,

                            localchecks, restricted, noinvariant):

        """

        select: the rql syntax tree Select node

        kwargs: query arguments

        localchecks: {(('Var name', (rqlexpr1, rqlexpr2)),

                       ('Var name1', (rqlexpr1, rqlexpr23))): [solution]}

              (see querier._check_permissions docstring for more information)

        restricted: set of variable names to which an rql expression has to be

              applied

        noinvariant: set of variable names that can't be considered has

              invariant due to security reason (will be filed by this method)

        """

        nbtrees = len(localchecks)

        myunion = union = select.parent

        # transform in subquery when len(localchecks)>1 and groups

        if nbtrees > 1 and (select.orderby or select.groupby or

                            select.having or select.has_aggregat or

                            select.distinct or

                            select.limit or select.offset):

            newselect = stmts.Select()

            # only select variables in subqueries

            origselection = select.selection

            select.select_only_variables()

            select.has_aggregat = False

            # create subquery first so correct node are used on copy

            # (eg ColumnAlias instead of Variable)

            aliases = [n.VariableRef(newselect.get_variable(vref.name, i))

                       for i, vref in enumerate(select.selection)]

            selected = set(vref.name for vref in aliases)

            # now copy original selection and groups

            for term in origselection:

                newselect.append_selected(term.copy(newselect))

            if select.orderby:

                sortterms = []

                for sortterm in select.orderby:

                    sortterms.append(sortterm.copy(newselect))

                    for fnode in sortterm.get_nodes(n.Function):

                        if fnode.name == 'FTIRANK':

                            # we've to fetch the has_text relation as well

                            var = fnode.children[0].variable

                            rel = iter(var.stinfo['ftirels']).next()

                            assert not rel.ored(), 'unsupported'

                            newselect.add_restriction(rel.copy(newselect))

                            # remove relation from the orig select and

                            # cleanup variable stinfo

                            rel.parent.remove(rel)

                            var.stinfo['ftirels'].remove(rel)

                            var.stinfo['relations'].remove(rel)

                            # XXX not properly re-annotated after security insertion?

                            newvar = newselect.get_variable(var.name)

                            newvar.stinfo.setdefault('ftirels', set()).add(rel)

                            newvar.stinfo.setdefault('relations', set()).add(rel)

                newselect.set_orderby(sortterms)

                _expand_selection(select.orderby, selected, aliases, select, newselect)

                select.orderby = () # XXX dereference?

            if select.groupby:

                newselect.set_groupby([g.copy(newselect) for g in select.groupby])

                _expand_selection(select.groupby, selected, aliases, select, newselect)

                select.groupby = () # XXX dereference?

            if select.having:

                newselect.set_having([g.copy(newselect) for g in select.having])

                _expand_selection(select.having, selected, aliases, select, newselect)

                select.having = () # XXX dereference?

            if select.limit:

                newselect.limit = select.limit

                select.limit = None

            if select.offset:

                newselect.offset = select.offset

                select.offset = 0

            myunion = stmts.Union()

            newselect.set_with([n.SubQuery(aliases, myunion)], check=False)

            newselect.distinct = select.distinct

            solutions = [sol.copy() for sol in select.solutions]

            cleanup_solutions(newselect, solutions)

            newselect.set_possible_types(solutions)

            # if some solutions doesn't need rewriting, insert original

            # select as first union subquery

            if () in localchecks:

                myunion.append(select)

            # we're done, replace original select by the new select with

            # subqueries (more added in the loop below)

            union.replace(select, newselect)

        elif not () in localchecks:

            union.remove(select)

        for lcheckdef, lchecksolutions in localchecks.iteritems():

            if not lcheckdef:

                continue

            myrqlst = select.copy(solutions=lchecksolutions)

            myunion.append(myrqlst)

            # in-place rewrite + annotation / simplification

            lcheckdef = [({var: 'X'}, rqlexprs) for var, rqlexprs in lcheckdef]

            self.rewrite(myrqlst, lcheckdef, kwargs)

            _add_noinvariant(noinvariant, restricted, myrqlst, nbtrees)

        if () in localchecks:

            select.set_possible_types(localchecks[()])

            add_types_restriction(self.schema, select)

            _add_noinvariant(noinvariant, restricted, select, nbtrees)

        self.annotate(union)

    def rewrite(self, select, snippets, kwargs, existingvars=None):

        """

        snippets: (varmap, list of rql expression)

                  with varmap a *tuple* (select var, snippet var)

        """

        self.select = select

        # remove_solutions used below require a copy

        self.solutions = solutions = select.solutions[:]

        self.kwargs = kwargs

        self.u_varname = None

        self.removing_ambiguity = False

        self.exists_snippet = {}

        self.pending_keys = []

        self.existingvars = existingvars

        # we have to annotate the rqlst before inserting snippets, even though

        # we'll have to redo it latter

        self.annotate(select)

        self.insert_snippets(snippets)

        if not self.exists_snippet and self.u_varname:

            # U has been inserted than cancelled, cleanup

            select.undefine_variable(select.defined_vars[self.u_varname])

        # clean solutions according to initial solutions

        newsolutions = remove_solutions(solutions, select.solutions,

                                        select.defined_vars)

        assert len(newsolutions) >= len(solutions), (

            'rewritten rql %s has lost some solutions, there is probably '

            'something wrong in your schema permission (for instance using a '

            'RQLExpression which insert a relation which doesn\'t exists in '

            'the schema)\nOrig solutions: %s\nnew solutions: %s' % (

            select, solutions, newsolutions))

        if len(newsolutions) > len(solutions):

            newsolutions = self.remove_ambiguities(snippets, newsolutions)

            assert newsolutions

        select.solutions = newsolutions

        add_types_restriction(self.schema, select)

    def insert_snippets(self, snippets, varexistsmap=None):

        self.rewritten = {}

        for varmap, rqlexprs in snippets:

            if isinstance(varmap, dict):

                varmap = tuple(sorted(varmap.items()))

            else:

                assert isinstance(varmap, tuple), varmap