"""RQL to SQL generator for native sources.

SQL queries optimization

~~~~~~~~~~~~~~~~~~~~~~~~

1. CWUser X WHERE X in_group G, G name 'users':

   CWUser is the only subject entity type for the in_group relation,

   which allow us to do ::

     SELECT eid_from FROM in_group, CWGroup

     WHERE in_group.eid_to = CWGroup.eid_from

     AND CWGroup.name = 'users'

2. Any X WHERE X nonfinal1 Y, Y nonfinal2 Z

   -> direct join between nonfinal1 and nonfinal2, whatever X,Y, Z (unless

      inlined...)

      NOT IMPLEMENTED (and quite hard to implement)

Potential optimization information is collected by the querier, sql generation

is done according to this information

: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"

import threading

from rql import BadRQLQuery, CoercionError

from rql.stmts import Union, Select

from rql.nodes import (SortTerm, VariableRef, Constant, Function, Not,

                       Variable, ColumnAlias, Relation, SubQuery, Exists)

from cubicweb import server

from cubicweb.server.sqlutils import SQL_PREFIX

from cubicweb.server.utils import cleanup_solutions

def _new_var(select, varname):

    newvar = select.get_variable(varname)

    if not 'relations' in newvar.stinfo:

        # not yet initialized

        newvar.prepare_annotation()

        newvar.stinfo['scope'] = select

        newvar._q_invariant = False

    return newvar

def _fill_to_wrap_rel(var, newselect, towrap, schema):

    for rel in var.stinfo['relations'] - var.stinfo['rhsrelations']:

        rschema = schema.rschema(rel.r_type)

        if rschema.inlined:

            towrap.add( (var, rel) )

            for vref in rel.children[1].iget_nodes(VariableRef):

                newivar = _new_var(newselect, vref.name)

                newselect.selection.append(VariableRef(newivar))

                _fill_to_wrap_rel(vref.variable, newselect, towrap, schema)

        elif rschema.final:

            towrap.add( (var, rel) )

def rewrite_unstable_outer_join(select, solutions, unstable, schema):

    """if some optional variables are unstable, they should be selected in a

    subquery. This function check this and rewrite the rql syntax tree if

    necessary (in place). Return a boolean telling if the tree has been modified

    """

    torewrite = set()

    modified = False

    for varname in tuple(unstable):

        var = select.defined_vars[varname]

        if not var.stinfo['optrelations']:

            continue

        modified = True

        unstable.remove(varname)

        torewrite.add(var)

        newselect = Select()

        newselect.need_distinct = newselect.need_intersect = False

        myunion = Union()

        myunion.append(newselect)

        # extract aliases / selection

        newvar = _new_var(newselect, var.name)

        newselect.selection = [VariableRef(newvar)]

        for avar in select.defined_vars.itervalues():

            if avar.stinfo['attrvar'] is var:

                newavar = _new_var(newselect, avar.name)

                newavar.stinfo['attrvar'] = newvar

                newselect.selection.append(VariableRef(newavar))

        towrap_rels = set()

        _fill_to_wrap_rel(var, newselect, towrap_rels, schema)

        # extract relations

        for var, rel in towrap_rels:

            newrel = rel.copy(newselect)

            newselect.add_restriction(newrel)

            select.remove_node(rel)

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

            newvar.stinfo['relations'].add(newrel)

            if rel.optional in ('left', 'both'):

                newvar.stinfo['optrelations'].add(newrel)

            for vref in newrel.children[1].iget_nodes(VariableRef):

                var = vref.variable

                var.stinfo['relations'].add(newrel)

                var.stinfo['rhsrelations'].add(newrel)

                if rel.optional in ('right', 'both'):

                    var.stinfo['optrelations'].add(newrel)

        # extract subquery solutions

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

        cleanup_solutions(newselect, mysolutions)

        newselect.set_possible_types(solutions)

        # full sub-query

        aliases = [VariableRef(select.get_variable(avar.name, i))

                   for i, avar in enumerate(newselect.selection)]

        select.add_subquery(SubQuery(aliases, myunion), check=False)

    return modified

def _new_solutions(rqlst, solutions):

    """first filter out subqueries variables from solutions"""

    newsolutions = []

    for origsol in solutions:

        asol = {}

        for vname in rqlst.defined_vars:

            asol[vname] = origsol[vname]

        if not asol in newsolutions:

            newsolutions.append(asol)

    return newsolutions

def remove_unused_solutions(rqlst, solutions, varmap, schema):

    """cleanup solutions: remove solutions where invariant variables are taking

    different types

    """

    newsols = _new_solutions(rqlst, solutions)

    existssols = {}

    unstable = set()

    invariants = {}

    for vname, var in rqlst.defined_vars.iteritems():

        vtype = newsols[0][vname]

        if var._q_invariant or vname in varmap:

            # remove invariant variable from solutions to remove duplicates

            # later, then reinserting a type for the variable even later

            for sol in newsols:

                invariants.setdefault(id(sol), {})[vname] = sol.pop(vname)

        elif var.scope is not rqlst:

            # move appart variables which are in a EXISTS scope and are variating

            try:

                thisexistssols, thisexistsvars = existssols[var.scope]

            except KeyError:

                thisexistssols = [newsols[0]]

                thisexistsvars = set()

                existssols[var.scope] = thisexistssols, thisexistsvars

            for i in xrange(len(newsols)-1, 0, -1):

                if vtype != newsols[i][vname]:

                    thisexistssols.append(newsols.pop(i))

                    thisexistsvars.add(vname)

        else:

            # remember unstable variables

            for i in xrange(1, len(newsols)):

                if vtype != newsols[i][vname]:

                    unstable.add(vname)

    if invariants:

        # filter out duplicates

        newsols_ = []

        for sol in newsols:

            if not sol in newsols_:

                newsols_.append(sol)

        newsols = newsols_

        # reinsert solutions for invariants

        for sol in newsols:

            for invvar, vartype in invariants[id(sol)].iteritems():

                sol[invvar] = vartype

        for sol in existssols:

            try:

                for invvar, vartype in invariants[id(sol)].iteritems():

                    sol[invvar] = vartype

            except KeyError:

                continue

    if len(newsols) > 1:

        if rewrite_unstable_outer_join(rqlst, newsols, unstable, schema):

            # remove variables extracted to subqueries from solutions

            newsols = _new_solutions(rqlst, newsols)

    return newsols, existssols, unstable

def relation_info(relation):

    lhs, rhs = relation.get_variable_parts()

    try:

        lhs = lhs.variable

        lhsconst = lhs.stinfo['constnode']

    except AttributeError:

        lhsconst = lhs

        lhs = None

    except KeyError:

        lhsconst = None # ColumnAlias

    try:

        rhs = rhs.variable

        rhsconst = rhs.stinfo['constnode']

    except AttributeError:

        rhsconst = rhs

        rhs = None

    except KeyError:

        rhsconst = None # ColumnAlias

    return lhs, lhsconst, rhs, rhsconst

def switch_relation_field(sql, table=''):

    switchedsql = sql.replace(table + '.eid_from', '__eid_from__')

    switchedsql = switchedsql.replace(table + '.eid_to',

                                      table + '.eid_from')

    return switchedsql.replace('__eid_from__', table + '.eid_to')

def sort_term_selection(sorts, selectedidx, rqlst, groups):

    # XXX beurk

    if isinstance(rqlst, list):

        def append(term):

            rqlst.append(term)

    else:

        def append(term):

            rqlst.selection.append(term.copy(rqlst))

    for sortterm in sorts:

        term = sortterm.term

        if not isinstance(term, Constant) and not str(term) in selectedidx:

            selectedidx.append(str(term))

            append(term)

            if groups:

                for vref in term.iget_nodes(VariableRef):

                    if not vref in groups:

                        groups.append(vref)

def fix_selection_and_group(rqlst, selectedidx, needwrap, selectsortterms,

                            sorts, groups, having):

    if selectsortterms and sorts:

        sort_term_selection(sorts, selectedidx, rqlst, not needwrap and groups)

    if sorts and groups:

        # when a query is grouped, ensure sort terms are grouped as well

        for sortterm in sorts:

            term = sortterm.term

            if not isinstance(term, Constant):

                for vref in term.iget_nodes(VariableRef):

                    if not vref in groups:

                        groups.append(vref)

    if needwrap:

        if groups:

            for vref in groups:

                if not vref.name in selectedidx:

                    selectedidx.append(vref.name)

                    rqlst.selection.append(vref)

        if having:

            for term in having:

                for vref in term.iget_nodes(VariableRef):

                    if not vref.name in selectedidx:

                        selectedidx.append(vref.name)

                        rqlst.selection.append(vref)

# IGenerator implementation for RQL->SQL ######################################

class StateInfo(object):

    def __init__(self, existssols, unstablevars):

        self.existssols = existssols

        self.unstablevars = unstablevars

        self.subtables = {}

    def reset(self, solution):

        """reset some visit variables"""

        self.solution = solution

        self.count = 0

        self.done = set()

        self.tables = self.subtables.copy()

        self.actual_tables = [[]]

        for _, tsql in self.tables.itervalues():

            self.actual_tables[-1].append(tsql)

        self.outer_tables = {}

        self.duplicate_switches = []

        self.aliases = {}

        self.restrictions = []

        self._restr_stack = []

        self.ignore_varmap = False

    def add_restriction(self, restr):

        if restr:

            self.restrictions.append(restr)

    def iter_exists_sols(self, exists):

        if not exists in self.existssols:

            yield 1

            return

        thisexistssols, thisexistsvars = self.existssols[exists]

        origsol = self.solution

        origtables = self.tables

        done = self.done

        for thisexistssol in thisexistssols:

            for vname in self.unstablevars:

                if thisexistssol[vname] != origsol[vname] and vname in thisexistsvars:

                    break

            else:

                self.tables = origtables.copy()

                self.solution = thisexistssol

                yield 1

                # cleanup self.done from stuff specific to exists

                for var in thisexistsvars:

                    if var in done:

                        done.remove(var)

                for rel in exists.iget_nodes(Relation):

                    if rel in done:

                        done.remove(rel)

        self.solution = origsol

        self.tables = origtables

    def push_scope(self):

        self.actual_tables.append([])

        self._restr_stack.append(self.restrictions)

        self.restrictions = []

    def pop_scope(self):

        restrictions = self.restrictions

        self.restrictions = self._restr_stack.pop()

        return restrictions, self.actual_tables.pop()

class SQLGenerator(object):

    """

    generation of SQL from the fully expanded RQL syntax tree

    SQL is designed to be used with a CubicWeb SQL schema

    Groups and sort are not handled here since they should not be handled at

    this level (see cubicweb.server.querier)

    we should not have errors here !

    WARNING: a CubicWebSQLGenerator instance is not thread safe, but generate is

    protected by a lock

    """

    def __init__(self, schema, dbms_helper, dbencoding='UTF-8', attrmap=None):

        self.schema = schema

        self.dbms_helper = dbms_helper

        self.dbencoding = dbencoding

        self.keyword_map = {'NOW' : self.dbms_helper.sql_current_timestamp,

                            'TODAY': self.dbms_helper.sql_current_date,

                            }

        if not self.dbms_helper.union_parentheses_support:

            self.union_sql = self.noparen_union_sql

        self._lock = threading.Lock()

        if attrmap is None:

            attrmap = {}

        self.attr_map = attrmap

    def generate(self, union, args=None, varmap=None):

        """return SQL queries and a variable dictionnary from a RQL syntax tree

        :partrqls: a list of couple (rqlst, solutions)

        :args: optional dictionary with values of substitutions used in the query

        :varmap: optional dictionary mapping variable name to a special table

          name, in case the query as to fetch data from temporary tables

        return an sql string and a dictionary with substitutions values

        """

        if args is None:

            args = {}

        if varmap is None:

            varmap =  {}

        self._lock.acquire()

        self._args = args

        self._varmap = varmap

        self._query_attrs = {}

        self._state = None

        self._not_scope_offset = 0

        try:

            # union query for each rqlst / solution

            sql = self.union_sql(union)

            # we are done

            return sql, self._query_attrs

        finally:

            self._lock.release()

    def union_sql(self, union, needalias=False): # pylint: disable-msg=E0202

        if len(union.children) == 1:

            return self.select_sql(union.children[0], needalias)

        sqls = ('(%s)' % self.select_sql(select, needalias)

                for select in union.children)

        return '\nUNION ALL\n'.join(sqls)

    def noparen_union_sql(self, union, needalias=False):

        # needed for sqlite backend which doesn't like parentheses around

        # union query. This may cause bug in some condition (sort in one of

        # the subquery) but will work in most case