"""plan execution of rql queries on a single source

:organization: Logilab

:copyright: 2001-2010 LOGILAB S.A. (Paris, FRANCE), license is LGPL v2.

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

:license: GNU Lesser General Public License, v2.1 - http://www.gnu.org/licenses

"""

__docformat__ = "restructuredtext en"

from copy import copy

from rql.stmts import Union, Select

from rql.nodes import Constant, Relation

from cubicweb import QueryError, typed_eid

from cubicweb.schema import VIRTUAL_RTYPES

from cubicweb.rqlrewrite import add_types_restriction

READ_ONLY_RTYPES = set(('eid', 'has_text', 'is', 'is_instance_of', 'identity'))

_CONSTANT = object()

_FROM_SUBSTEP = object()

def _extract_const_attributes(plan, rqlst, to_build):

    """add constant values to entity def, mark variables to be selected

    """

    to_select = {}

    for relation in rqlst.main_relations:

        lhs, rhs = relation.get_variable_parts()

        rtype = relation.r_type

        if rtype in READ_ONLY_RTYPES:

            raise QueryError("can't assign to %s" % rtype)

        try:

            edef = to_build[str(lhs)]

        except KeyError:

            # lhs var is not to build, should be selected and added as an

            # object relation

            edef = to_build[str(rhs)]

            to_select.setdefault(edef, []).append((rtype, lhs, 1))

        else:

            if isinstance(rhs, Constant) and not rhs.uid:

                # add constant values to entity def

                value = rhs.eval(plan.args)

                eschema = edef.e_schema

                attrtype = eschema.subjrels[rtype].objects(eschema)[0]

                if attrtype == 'Password' and isinstance(value, unicode):

                    value = value.encode('UTF8')

                edef[rtype] = value

            elif to_build.has_key(str(rhs)):

                # create a relation between two newly created variables

                plan.add_relation_def((edef, rtype, to_build[rhs.name]))

            else:

                to_select.setdefault(edef, []).append( (rtype, rhs, 0) )

    return to_select

def _extract_eid_consts(plan, rqlst):

    """return a dict mapping rqlst variable object to their eid if specified in

    the syntax tree

    """

    session = plan.session

    eschema = session.vreg.schema.eschema

    if rqlst.where is None:

        return {}

    eidconsts = {}

    neweids = session.transaction_data.get('neweids', ())

    for rel in rqlst.where.get_nodes(Relation):

        if rel.r_type == 'eid' and not rel.neged(strict=True):

            lhs, rhs = rel.get_variable_parts()

            if isinstance(rhs, Constant):

                eid = typed_eid(rhs.eval(plan.args))

                # check read permission here since it may not be done by

                # the generated select substep if not emited (eg nothing

                # to be selected)

                if eid not in neweids:

                    eschema(session.describe(eid)[0]).check_perm(session, 'read')

                eidconsts[lhs.variable] = eid

    return eidconsts

class SSPlanner(object):

    """SingleSourcePlanner: build execution plan for rql queries

    optimized for single source repositories

    """

    def __init__(self, schema, rqlhelper):

        self.schema = schema

        self.rqlhelper = rqlhelper

    def build_plan(self, plan):

        """build an execution plan from a RQL query

        do nothing here, dispatch according to the statement type

        """

        build_plan = getattr(self, 'build_%s_plan' % plan.rqlst.TYPE)

        for step in build_plan(plan, plan.rqlst):

            plan.add_step(step)

    def build_select_plan(self, plan, rqlst):

        """build execution plan for a SELECT RQL query. Suppose only one source

        is available and so avoid work need for query decomposition among sources

        the rqlst should not be tagged at this point.

        """

        plan.preprocess(rqlst)

        return (OneFetchStep(plan, rqlst, plan.session.repo.sources),)

    def build_insert_plan(self, plan, rqlst):

        """get an execution plan from an INSERT RQL query"""

        # each variable in main variables is a new entity to insert

        to_build = {}

        session = plan.session

        etype_class = session.vreg['etypes'].etype_class

        for etype, var in rqlst.main_variables:

            # need to do this since entity class is shared w. web client code !

            to_build[var.name] = etype_class(etype)(session)

            plan.add_entity_def(to_build[var.name])

        # add constant values to entity def, mark variables to be selected

        to_select = _extract_const_attributes(plan, rqlst, to_build)

        # add necessary steps to add relations and update attributes

        step = InsertStep(plan) # insert each entity and its relations

        step.children += self._compute_relation_steps(plan, rqlst, to_select)

        return (step,)

    def _compute_relation_steps(self, plan, rqlst, to_select):

        """handle the selection of relations for an insert query"""

        eidconsts = _extract_eid_consts(plan, rqlst)

        for edef, rdefs in to_select.items():

            # create a select rql st to fetch needed data

            select = Select()

            eschema = edef.e_schema

            for i, (rtype, term, reverse) in enumerate(rdefs):

                if getattr(term, 'variable', None) in eidconsts:

                    value = eidconsts[term.variable]

                else:

                    select.append_selected(term.copy(select))

                    value = _FROM_SUBSTEP

                if reverse:

                    rdefs[i] = (rtype, InsertRelationsStep.REVERSE_RELATION, value)

                else:

                    rschema = eschema.subjrels[rtype]

                    if rschema.final or rschema.inlined:

                        rdefs[i] = (rtype, InsertRelationsStep.FINAL, value)

                    else:

                        rdefs[i] = (rtype, InsertRelationsStep.RELATION, value)

            step = InsertRelationsStep(plan, edef, rdefs)

                step.children += self._select_plan(plan, select, rqlst.solutions)

            yield step

    def build_delete_plan(self, plan, rqlst):

        """get an execution plan from a DELETE RQL query"""

        # build a select query to fetch entities to delete

        steps = []

        for etype, var in rqlst.main_variables:

            step = DeleteEntitiesStep(plan)

            step.children += self._sel_variable_step(plan, rqlst.solutions,

                                                     rqlst.where, etype, var)

            steps.append(step)

        for relation in rqlst.main_relations:

            step = DeleteRelationsStep(plan, relation.r_type)

            step.children += self._sel_relation_steps(plan, rqlst.solutions,

                                                      rqlst.where, relation)

            steps.append(step)

        return steps

    def _sel_variable_step(self, plan, solutions, restriction, etype, varref):

        """handle the selection of variables for a delete query"""

        select = Select()

        varref = varref.copy(select)

        select.defined_vars = {varref.name: varref.variable}

        select.append_selected(varref)

        if restriction is not None:

            select.set_where(restriction.copy(select))

        if etype != 'Any':

            select.add_type_restriction(varref.variable, etype)

        return self._select_plan(plan, select, solutions)

    def _sel_relation_steps(self, plan, solutions, restriction, relation):

        """handle the selection of relations for a delete query"""

        select = Select()

        lhs, rhs = relation.get_variable_parts()

        select.append_selected(lhs.copy(select))

        select.append_selected(rhs.copy(select))

        select.set_where(relation.copy(select))

        if restriction is not None:

            select.add_restriction(restriction.copy(select))

        return self._select_plan(plan, select, solutions)

    def build_set_plan(self, plan, rqlst):

        """get an execution plan from an SET RQL query"""

        getrschema = self.schema.rschema

        select = Select()   # potential substep query

        selectedidx = {}    # local state

        attributes = set()  # edited attributes

        updatedefs = []     # definition of update attributes/relations

        selidx = residx = 0 # substep selection / resulting rset indexes

        # search for eid const in the WHERE clause

        eidconsts = _extract_eid_consts(plan, rqlst)

        # build `updatedefs` describing things to update and add necessary

        # variables to the substep selection

        for i, relation in enumerate(rqlst.main_relations):

            if relation.r_type in VIRTUAL_RTYPES:

                raise QueryError('can not assign to %r relation'

                                 % relation.r_type)

            lhs, rhs = relation.get_variable_parts()

            lhskey = lhs.as_string('utf-8')

            if not lhskey in selectedidx:

                if lhs.variable in eidconsts:

                    eid = eidconsts[lhs.variable]

                    lhsinfo = (_CONSTANT, eid, residx)

                else:

                    select.append_selected(lhs.copy(select))

                    lhsinfo = (_FROM_SUBSTEP, selidx, residx)

                    selidx += 1

                residx += 1

                selectedidx[lhskey] = lhsinfo

            else:

                lhsinfo = selectedidx[lhskey][:-1] + (None,)

            rhskey = rhs.as_string('utf-8')

            if not rhskey in selectedidx:

                if isinstance(rhs, Constant):

                    rhsinfo = (_CONSTANT, rhs.eval(plan.args), residx)

                elif getattr(rhs, 'variable', None) in eidconsts:

                    eid = eidconsts[rhs.variable]

                    rhsinfo = (_CONSTANT, eid, residx)

                else:

                    select.append_selected(rhs.copy(select))

                    rhsinfo = (_FROM_SUBSTEP, selidx, residx)

                    selidx += 1

                residx += 1

                selectedidx[rhskey] = rhsinfo

            else:

                rhsinfo = selectedidx[rhskey][:-1] + (None,)

            rschema = getrschema(relation.r_type)

            updatedefs.append( (lhsinfo, rhsinfo, rschema) )

            if rschema.final or rschema.inlined:

                attributes.add(relation.r_type)

        # the update step

        step = UpdateStep(plan, updatedefs, attributes)

        # when necessary add substep to fetch yet unknown values

            # set distinct to avoid potential duplicate key error

            select.distinct = True

            step.children += self._select_plan(plan, select, rqlst.solutions)

        return (step,)

    # internal methods ########################################################

    def _select_plan(self, plan, select, solutions):

        union = Union()

        union.append(select)

        select.clean_solutions(solutions)

        add_types_restriction(self.schema, select)

        self.rqlhelper.annotate(union)

        return self.build_select_plan(plan, union)

# execution steps and helper functions ########################################

def varmap_test_repr(varmap, tablesinorder):

    if varmap is None:

        return varmap

    maprepr = {}

    for var, sql in varmap.iteritems():

        table, col = sql.split('.')

        maprepr[var] = '%s.%s' % (tablesinorder[table], col)

    return maprepr

def offset_result(offset, result):

    offset -= len(result)

    if offset < 0:

        result = result[offset:]

        offset = None

    elif offset == 0:

        offset = None

        result = ()

    return offset, result

class LimitOffsetMixIn(object):

    limit = offset = None

    def set_limit_offset(self, limit, offset):

        self.limit = limit

        self.offset = offset or None

class Step(object):

    """base abstract class for execution step"""

    def __init__(self, plan):

        self.plan = plan

        self.children = []

    def execute_child(self):

        assert len(self.children) == 1

        return self.children[0].execute()

    def execute_children(self):

        for step in self.children:

            step.execute()

    def execute(self):

        """execute this step and store partial (eg this step) results"""

        raise NotImplementedError()

    def mytest_repr(self):

        """return a representation of this step suitable for test"""

        return (self.__class__.__name__,)

    def test_repr(self):

        """return a representation of this step suitable for test"""

        return self.mytest_repr() + (

            [step.test_repr() for step in self.children],)

class OneFetchStep(LimitOffsetMixIn, Step):

    """step consisting in fetching data from sources and directly returning

    results

    """

    def __init__(self, plan, union, sources, inputmap=None):

        Step.__init__(self, plan)

        self.union = union

        self.sources = sources

        self.inputmap = inputmap

        self.set_limit_offset(union.children[-1].limit, union.children[-1].offset)

    def set_limit_offset(self, limit, offset):

        LimitOffsetMixIn.set_limit_offset(self, limit, offset)

        for select in self.union.children:

            select.limit = limit

            select.offset = offset

    def execute(self):

        """call .syntax_tree_search with the given syntax tree on each

        source for each solution

        """

        self.execute_children()

        session = self.plan.session

        args = self.plan.args

        inputmap = self.inputmap