# copyright 2003-2014 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.

"""

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.items():

            # 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 sorted(possibletypes):

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

                else:

                    etype = next(iter(possibletypes))

                    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 cst.value not 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 vref.name not 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 _compatible_relation(relations, stmt, sniprel):

    """Search among given rql relation nodes if there is one 'compatible' with the

    snippet relation, and return it if any, else None.

    A relation is compatible if it:

    * belongs to the currently processed statement,

    * isn't negged (i.e. direct parent is a NOT node)

    * isn't optional (outer join) or similarly as the snippet relation

    """

    for rel in relations:

        # don't share if relation's scope is not the current statement

        if rel.scope is not stmt:

            continue

        # don't share neged relation

        if rel.neged(strict=True):

            continue

        # don't share optional relation, unless the snippet relation is

        # similarly optional

        if rel.optional and rel.optional != sniprel.optional:

            continue

        return rel

    return None

def iter_relations(stinfo):

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

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

def need_exists(node):

    """Return true if the given node should be wrapped in an `Exists` node.

    This is true when node isn't already an `Exists` or `Not` node, nor a

    `And`/`Or` of `Exists` or `Not` nodes.

    """

    if isinstance(node, (n.Exists, n.Not)):

        return False

    if isinstance(node, (n.Or, n.And)):

        return need_exists(node.children[0]) or need_exists(node.children[1])

    return True

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 VariableFromSubQuery(Exception):

    """flow control exception to indicate that a variable is coming from a

    subquery, and let parent act accordingly

    """

    def __init__(self, variable):

        self.variable = variable

class RQLRewriter(object):

    """Insert some rql snippets into another rql syntax tree, for security /

    relation vocabulary. This implies that it should only restrict results of

    the original query, not generate new ones. Hence, inserted snippets are

    inserted under an EXISTS node.

    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.compute_var_types

    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 = next(iter(var.stinfo['ftirels']))

                            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.items():

            if not lcheckdef:

                continue

            myrqlst = select.copy(solutions=lchecksolutions)

            myunion.append(myrqlst)

            # in-place rewrite + annotation / simplification

            lcheckdef = [({v: 'X'}, rqlexprs) for v, 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 *dict* {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 later

        self.annotate(select)

        self.insert_snippets(snippets)

        if not self.exists_snippet and self.u_varname: