"""RQL to SQL generator for native sources.

SQL queries optimization

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

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

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

   which allow us to do ::

     SELECT eid_from FROM in_group, EGroup

     WHERE in_group.eid_to = EGroup.eid_from

     AND EGroup.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-2008 LOGILAB S.A. (Paris, FRANCE), all rights reserved.

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

"""

__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.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.is_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

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

        cleanup_solutions(newselect, solutions)

        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

    """

    newsolutions = _new_solutions(rqlst, solutions)

    existssols = {}

    unstable = set()

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

        vtype = newsolutions[0][vname]

        if var._q_invariant or vname in varmap:

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

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

                    newsolutions.pop(i)

        elif not var.scope is rqlst:

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

            try:

                thisexistssols, thisexistsvars = existssols[var.scope]

            except KeyError:

                thisexistssols = [newsolutions[0]]

                thisexistsvars = set()

                existssols[var.scope] = thisexistssols, thisexistsvars

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

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

                    thisexistssols.append(newsolutions.pop(i))

                    thisexistsvars.add(vname)

        else:

            # remember unstable variables

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

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

                    unstable.add(vname)

    if len(newsolutions) > 1:

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

            # remove variables extracted to subqueries from solutions

            newsolutions = _new_solutions(rqlst, newsolutions)

    return newsolutions, 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(rqlst, selectedidx, needwrap, sorts, groups, having):

    if sorts:

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

    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.attr_vars = {}

        self.aliases = {}

        self.restrictions = []

        self._restr_stack = []

    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'):

        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()

    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

        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

        # see http://www.sqlite.org/cvstrac/tktview?tn=3074

        sqls = (self.select_sql(select, needalias)

                for i, select in enumerate(union.children))

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

    def select_sql(self, select, needalias=False):

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

        :select: a selection statement of the syntax tree (`rql.stmts.Select`)

        :solution: a dictionnary containing variables binding.

          A solution's dictionnary has variable's names as key and variable's

          types as values

        :needwrap: boolean telling if the query will be wrapped in an outer

          query (to deal with aggregat and/or grouping)

        """

        distinct = selectsortterms = select.need_distinct

        sorts = select.orderby

        groups = select.groupby

        having = select.having

        # remember selection, it may be changed and have to be restored

        origselection = select.selection[:]

        # check if the query will have union subquery, if it need sort term

        # selection (union or distinct query) and wrapping (union with groups)

        needwrap = False

        sols = select.solutions

        if len(sols) > 1:

            # remove invariant from solutions

            sols, existssols, unstable = remove_unused_solutions(

                select, sols, self._varmap, self.schema)

            if len(sols) > 1:

                # if there is still more than one solution, a UNION will be

                # generated and so sort terms have to be selected

                selectsortterms = True

                # and if select is using group by or aggregat, a wrapping

                # query will be necessary

                if groups or select.has_aggregat:

                    select.select_only_variables()

                    needwrap = True                        

        else:

            existssols, unstable = {}, ()

        state = StateInfo(existssols, unstable)

        # treat subqueries

        self._subqueries_sql(select, state)