# copyright 2003-2011 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/>.
"""Defines the diferent querier steps usable in plans.

FIXME : this code needs refactoring. Some problems :
* get data from the parent plan, the latest step, temporary table...
* each step has is own members (this is not necessarily bad, but a bit messy
  for now)
"""
from __future__ import with_statement

__docformat__ = "restructuredtext en"

from rql.nodes import VariableRef, Variable, Function

from cubicweb.server.ssplanner import (LimitOffsetMixIn, Step, OneFetchStep,
                                    varmap_test_repr, offset_result)

AGGR_TRANSFORMS = {'COUNT':'SUM', 'MIN':'MIN', 'MAX':'MAX', 'SUM': 'SUM'}

class remove_and_restore_clauses(object):
    def __init__(self, union, keepgroup):
        self.union = union
        self.keepgroup = keepgroup
        self.clauses = None

    def __enter__(self):
        self.clauses = clauses = []
        for select in self.union.children:
            if self.keepgroup:
                having, orderby = select.having, select.orderby
                select.having, select.orderby = (), ()
                clauses.append( (having, orderby) )
            else:
                groupby, having, orderby = select.groupby, select.having, select.orderby
                select.groupby, select.having, select.orderby = (), (), ()
                clauses.append( (groupby, having, orderby) )

    def __exit__(self, exctype, exc, traceback):
        for i, select in enumerate(self.union.children):
            if self.keepgroup:
                select.having, select.orderby = self.clauses[i]
            else:
                select.groupby, select.having, select.orderby = self.clauses[i]


class FetchStep(OneFetchStep):
    """step consisting in fetching data from sources, and storing result in
    a temporary table
    """
    def __init__(self, plan, union, sources, table, keepgroup, inputmap=None):
        OneFetchStep.__init__(self, plan, union, sources)
        # temporary table to store step result
        self.table = table
        # should groupby clause be kept or not
        self.keepgroup = keepgroup
        # variables mapping to use as input
        self.inputmap = inputmap
        # output variable mapping
        srqlst = union.children[0] # sample select node
        # add additional information to the output mapping
        self.outputmap = plan.init_temp_table(table, srqlst.selection,
                                              srqlst.solutions[0])
        for vref in srqlst.selection:
            if not isinstance(vref, VariableRef):
                continue
            var = vref.variable
            if var.stinfo.get('attrvars'):
                for lhsvar, rtype in var.stinfo['attrvars']:
                    if lhsvar.name in srqlst.defined_vars:
                        key = '%s.%s' % (lhsvar.name, rtype)
                        self.outputmap[key] = self.outputmap[var.name]
            else:
                rschema = self.plan.schema.rschema
                for rel in var.stinfo['rhsrelations']:
                    if rschema(rel.r_type).inlined:
                        lhsvar = rel.children[0]
                        if lhsvar.name in srqlst.defined_vars:
                            key = '%s.%s' % (lhsvar.name, rel.r_type)
                            self.outputmap[key] = self.outputmap[var.name]

    def execute(self):
        """execute this step"""
        self.execute_children()
        plan = self.plan
        plan.create_temp_table(self.table)
        union = self.union
        with remove_and_restore_clauses(union, self.keepgroup):
            for source in self.sources:
                source.flying_insert(self.table, plan.session, union, plan.args,
                                     self.inputmap)

    def mytest_repr(self):
        """return a representation of this step suitable for test"""
        with remove_and_restore_clauses(self.union, self.keepgroup):
            try:
                inputmap = varmap_test_repr(self.inputmap, self.plan.tablesinorder)
                outputmap = varmap_test_repr(self.outputmap, self.plan.tablesinorder)
            except AttributeError:
                inputmap = self.inputmap
                outputmap = self.outputmap
            return (self.__class__.__name__,
                    sorted((r.as_string(kwargs=self.plan.args), r.solutions)
                           for r in self.union.children),
                    sorted(self.sources), inputmap, outputmap)


class AggrStep(LimitOffsetMixIn, Step):
    """step consisting in making aggregat from temporary data in the system
    source
    """
    def __init__(self, plan, selection, select, table, outputtable=None):
        Step.__init__(self, plan)
        # original selection
        self.selection = selection
        # original Select RQL tree
        self.select = select
        # table where are located temporary results
        self.table = table
        # optional table where to write results
        self.outputtable = outputtable
        if outputtable is not None:
            plan.init_temp_table(outputtable, selection, select.solutions[0])

        #self.inputmap = inputmap

    def mytest_repr(self):
        """return a representation of this step suitable for test"""
        try:
            # rely on a monkey patch (cf unittest_querier)
            table = self.plan.tablesinorder[self.table]
            outputtable = self.outputtable and self.plan.tablesinorder[self.outputtable]
        except AttributeError:
            # not monkey patched
            table = self.table
            outputtable = self.outputtable
        sql = self.get_sql().replace(self.table, table)
        return (self.__class__.__name__, sql, outputtable)

    def execute(self):
        """execute this step"""
        self.execute_children()
        sql = self.get_sql()
        if self.outputtable:
            self.plan.create_temp_table(self.outputtable)
            sql = 'INSERT INTO %s %s' % (self.outputtable, sql)
        return self.plan.sqlexec(sql, self.plan.args)

    def get_sql(self):
        self.inputmap = inputmap = self.children[-1].outputmap
        dbhelper=self.plan.syssource.dbhelper
        # get the select clause
        clause = []
        for i, term in enumerate(self.selection):
            try:
                var_name = inputmap[term.as_string()]
            except KeyError:
                var_name = 'C%s' % i
            if isinstance(term, Function):
                # we have to translate some aggregat function
                # (for instance COUNT -> SUM)
                orig_name = term.name
                try:
                    term.name = AGGR_TRANSFORMS[term.name]
                    # backup and reduce children
                    orig_children = term.children
                    term.children = [VariableRef(Variable(var_name))]
                    clause.append(term.accept(self))
                    # restaure the tree XXX necessary?
                    term.name = orig_name
                    term.children = orig_children
                except KeyError:
                    clause.append(var_name)
            else:
                clause.append(var_name)
                for vref in term.iget_nodes(VariableRef):
                    inputmap[vref.name] = var_name
        # XXX handle distinct with non selected sort term
        if self.select.distinct:
            sql = ['SELECT DISTINCT %s' % ', '.join(clause)]
        else:
            sql = ['SELECT %s' % ', '.join(clause)]
        sql.append("FROM %s" % self.table)
        # get the group/having clauses
        if self.select.groupby:
            clause = [inputmap[var.name] for var in self.select.groupby]
            grouped = set(var.name for var in self.select.groupby)
            sql.append('GROUP BY %s' % ', '.join(clause))
        else:
            grouped = None
        if self.select.having:
            clause = [term.accept(self) for term in self.select.having]
            sql.append('HAVING %s' % ', '.join(clause))
        # get the orderby clause
        if self.select.orderby:
            clause = []
            for sortterm in self.select.orderby:
                sqlterm = sortterm.term.accept(self)
                if sortterm.asc:
                    clause.append(sqlterm)
                else:
                    clause.append('%s DESC' % sqlterm)
                if grouped is not None:
                    for vref in sortterm.iget_nodes(VariableRef):
                        if not vref.name in grouped:
                            sql[-1] += ', ' + self.inputmap[vref.name]
                            grouped.add(vref.name)
            sql = dbhelper.sql_add_order_by(' '.join(sql),
                                            clause,
                                            None, False,
                                            self.limit or self.offset)
        else:
            sql = ' '.join(sql)
            clause = None

        sql = dbhelper.sql_add_limit_offset(sql, self.limit, self.offset, clause)
        return sql

    def visit_function(self, function):
        """generate SQL name for a function"""
        try:
            return self.children[0].outputmap[str(function)]
        except KeyError:
            return '%s(%s)' % (function.name,
                               ','.join(c.accept(self) for c in function.children))

    def visit_variableref(self, variableref):
        """get the sql name for a variable reference"""
        try:
            return self.inputmap[variableref.name]
        except KeyError: # XXX duh? explain
            return variableref.variable.name

    def visit_constant(self, constant):
        """generate SQL name for a constant"""
        assert constant.type == 'Int'
        return str(constant.value)


class UnionStep(LimitOffsetMixIn, Step):
    """union results of child in-memory steps (e.g. OneFetchStep / AggrStep)"""

    def execute(self):
        """execute this step"""
        result = []
        limit = olimit = self.limit
        offset = self.offset
        assert offset != 0
        if offset is not None:
            limit = limit + offset
        for step in self.children:
            if limit is not None:
                if offset is None:
                    limit = olimit - len(result)
                step.set_limit_offset(limit, None)
            result_ = step.execute()
            if offset is not None:
                offset, result_ = offset_result(offset, result_)
            result += result_
            if limit is not None:
                if len(result) >= olimit:
                    return result[:olimit]
        return result

    def mytest_repr(self):
        """return a representation of this step suitable for test"""
        return (self.__class__.__name__, self.limit, self.offset)


class IntersectStep(UnionStep):
    """return intersection of results of child in-memory steps (e.g. OneFetchStep / AggrStep)"""

    def execute(self):
        """execute this step"""
        result = set()
        for step in self.children:
            result &= frozenset(step.execute())
        result = list(result)
        if self.offset:
            result = result[self.offset:]
        if self.limit:
            result = result[:self.limit]
        return result


class UnionFetchStep(Step):
    """union results of child steps using temporary tables (e.g. FetchStep)"""

    def execute(self):
        """execute this step"""
        self.execute_children()


__all__ = ('FetchStep', 'AggrStep', 'UnionStep', 'UnionFetchStep', 'IntersectStep')