# Copyright 2011 Viewfinder Inc. All Rights Reserved.

"""Parses query expressions into hierarchical query trees.

Given a data schema (defined via schema.py), you can create Query objects
and execute them against a database client. For example, with table 'User'
and columns 'given_name' and 'Email', you can run a query as follows:

  query_str = '(user.given_name=spencer | user.given_name="andrew") & user.family_name=kimball'
  users = User.Query(query_str).Evaluate(client)


  Query: query object to evaluate query expression.

Supports parameterized phrases which should be used with strings from an untrusted source in order to mitigate
injection attacks.  The keys are valid in the phrase part of the query where the key name is surrounded by braces.
Parameterized queries are passed as a tuple where the first element is the query string followed by a dictionary
with the parameter keys mapped to parameter values.  The above query expression using parameters:

  bound_query_str = ('(user.given_name={p1} | user.given_name={p2}) & user.family_name={p3}',
                     {'p1': 'spencer', 'p2': 'andrew', 'p3': 'kimball'})

The parameter keys may be made up of letters, digits, and the underscore (_).

Based on calculator.py by Andrew Brehaut & Steven Ashley of picoparse.
"""

__author__ = 'spencer@emailscrubbed.com (Spencer Kimball)'

import logging
import re

from bisect import bisect_left, bisect_right
from collections import namedtuple
from functools import partial
from picoparse import one_of, choice, many1, tri, commit, p, many_until1
from picoparse.text import run_text_parser, as_string, lexeme, whitespace, quoted
from string import digits, letters
from tornado import escape
from viewfinder.backend.base import util
from viewfinder.backend.db import db_client, vf_schema
from viewfinder.backend.db.schema import IndexedTable

_query_cache = util.LRUCache(100)

# Optimization for the common case:  Indexer.GetQueryString usually returns a single token
# as a quoted string.  Picoparse's character-by-character operation is slow, so bypass
# it when we can.
_QUOTED_TOKEN_RE = re.compile(r'^"([^"\\]*)"$')

class _MatchResult(object):
  """The result of a successful match after node evaluation. Data
  initially contains the value attached to a key in a posting list,
  but may be arbitrarily updated and augmented as it moves through the
  query tree.
  """
  __slots__ = ['key', 'data']

  def __init__(self, key, data):
    self.key = key
    self.data = data

  def __cmp__(self, other):
    return cmp(self.key, other.key)


# The result from evaluating a node or subtree of the query tree.
# Matches is a list of '_MatchResult' values. 'last_key' is the
# last key in the range which was successfully evaluated.
_EvalResult = namedtuple('_EvalResult', ['matches', 'last_key', 'read_units'])


class _QueryNode(object):
  """Query nodes form binary trees with set operations at each
  intermediate node and ranges of keys at leaf nodes.
  """
  pass


class _OpNode(_QueryNode):
  """Operation nodes implement set operations on the results of
  key range queries. Set operations include union, difference and
  intersection.
  """
  def __init__(self, schema, left):
    """Takes the left child as a parameter. The right child is merged
    into this operation via a call to Merge. The two are followed
    recursively during evaluation, with the left and right evaluated
    asynchronously. When both have completed, the result of the set
    operation (union, difference, intersection & positional intersection)
    is evaluated and returned to parent.
    """
    self._left = left
    self._right = None

  def PrintTree(self, level, param_dict):
    """Depth-first printout of tree structure for debugging."""
    return self._OpName() + ' __ ' + self._left.PrintTree(level + 1, param_dict) + '\n' + level * ('     ') + \
                            '  \_ ' + self._right.PrintTree(level + 1, param_dict)

  def Merge(self, right):
    """Rearranges the parent-child relationship to conform to operator
    precedence. If the precedence of the current node is >= the node
    to its right, then the right-hand node becomes the new parent,
    with the current node made its new left-hand node (meaning this
    node will be evaluated first, as the evaluation is a depth-first
    traversal). Otherwise, the right-hand node simply stays the right
    hand node.
    """
    if self._precedence >= right._precedence:
      self._right = right._left
      right._left = self
      return right
    else:
      self._right = right
      return self

  def Evaluate(self, client, callback, start_key, consistent_read, param_dict):
    """Recursively evaluates the query tree via a depth- first
    traversal. Returns the result set, defined by the data delivered
    via the IndexTermNodes and then operated on by the OpNodes.
    """
    with util.ArrayBarrier(partial(self._SetOperation, callback)) as b:
      self._left.Evaluate(client, b.Callback(), start_key, consistent_read, param_dict)
      self._right.Evaluate(client, b.Callback(), start_key, consistent_read, param_dict)

  def _OpName(self):
    raise NotImplementedError()

  def _SetOperation(self):
    raise NotImplementedError()


class Union(_OpNode):
  """Returns the union of two sets. Lowest precedence."""
  def __init__(self, schema, left):
    super(Union, self).__init__(schema, left)
    self._precedence = 0

  def _OpName(self):
    return "| "

  def _SetOperation(self, callback, results):
    """For union, the sets are additively combined. Both the first and
    last keys are defined as the minimum of first and last keys for
    the two results. This jibes well with the intuitive notion that we
    were able to successfully evaluate the union between these two
    sets for all values starting at the very minimum up to and
    including the minimum of the two last keys. This prevents us from
    over-stepping a gap between the two ranges for example.
    """
    assert len(results) == 2
    matches = results[0].matches + results[1].matches
    matches.sort()

    if results[0].last_key is None:
      last_key = results[1].last_key
    elif results[1].last_key is None:
      last_key = results[0].last_key
    else:
      last_key = min(results[0].last_key, results[1].last_key)

    # Only returns matches up to last_key.
    if last_key:
      matches = matches[:bisect_right(matches, _MatchResult(last_key, None))]
    callback(_EvalResult(matches=matches, last_key=last_key,
                         read_units=results[0].read_units + results[1].read_units))


class Difference(_OpNode):
  """Returns the difference of two sets."""
  def __init__(self, schema, left):
    super(Difference, self).__init__(schema, left)
    self._precedence = 1

  def _OpName(self):
    return "- "

  def _SetOperation(self, callback, results):
    """For set difference, the ordering matters. The second set can be
    thought of as a mask over the first set. The overlap starting at
    the first key and extending to the min(last_key1, last_key2) defines
    the successfully-evaluated range. To see this, consider that any
    portion of the subtracted range which lies before the first set is
    irrelevant (as we can only eval what we have of the first range
    going forward), and any portion of the subtracted range which lies
    after the first set can't yet be evaluated, as we don't yet know
    the remainder of the first set.
    """
    assert len(results) == 2
    m1 = results[0].matches
    m2 = results[1].matches
    matches = []
    while m1:
      m2_idx = bisect_left(m2, m1[0])
      if m2_idx == len(m2):
        break
      elif m1[0] == m2[m2_idx]:
        # A match, so skip.
        m1 = m1[1:]
      elif m2_idx == 0:
        # No match, so include the range of m1's up to the start of m2.
        m1_idx = bisect_left(m1, m2[0])
        if m1_idx == len(m1):
          # The last result of m1 is still before the next m2; include all.
          matches += m1
          break
        else:
          # Include up to the next possible match between m1 and m2.
          matches += m1[:m1_idx]
          m1 = m1[m1_idx:]

    if results[0].last_key is None:
      last_key = results[1].last_key
    elif results[1].last_key is None:
      last_key = results[0].last_key
    else:
      last_key = min(results[0].last_key, results[1].last_key)
    callback(_EvalResult(matches=matches, last_key=last_key,
                         read_units=results[0].read_units + results[1].read_units))


class Intersection(_OpNode):
  """Returns the intersection of two sets."""
  def __init__(self, schema, left):
    super(Intersection, self).__init__(schema, left)
    self._precedence = 2

  def _OpName(self):
    return "& "

  def _SetOperation(self, callback, results):
    """Efficiently skip through the two lists using list bisection.
    """
    assert len(results) == 2
    m1 = results[0].matches
    m2 = results[1].matches
    matches = []
    while m1 and m2:
      if m1[0] < m2[0]:
        m1 = m1[bisect_left(m1, m2[0]):]
      elif m2[0] < m1[0]:
        m2 = m2[bisect_left(m2, m1[0]):]
      else:
        matches.append(m1[0])
        m1 = m1[1:]
        m2 = m2[1:]

    callback(_EvalResult(matches=matches, last_key=self._ComputeLastKey(results),
                         read_units=results[0].read_units + results[1].read_units))

  def _ComputeLastKey(self, results):
    """For set intersection, the successfully-evaluated portion starts
    at the maximum of the first keys and extends as far as the minimum
    of the last keys. This is the portion for which we have enough
    information to definitively determine intersection.
    """
    if results[0].last_key is None:
      last_key = results[1].last_key
    elif results[1].last_key is None:
      last_key = results[0].last_key
    else:
      last_key = min(results[0].last_key, results[1].last_key)
    return last_key


class PositionalIntersection(Intersection):
  """Returns the intersection of two posting lists, but only return
  matches where relative positions between the left and right nodes
  are self._delta apart.
  """
  def __init__(self, schema, left):
    super(PositionalIntersection, self).__init__(schema, left)
    self._delta = 1

  def Merge(self, right):
    """Merging for positional intersections works a little differently
    than for normal operations. Here, we are on the lookout to merge
    place-holder (None) nodes out of the tree. If the right-hand node
    we're trying to merge is a place-holder, then return the left-hand
    node as the result of the merge (this removes trailing
    place-holders). If the right-hand node we're merging's left-hand
    node is a place-holder, we skip over the right-hand node by
    attaching its right-hand node in its place. This operation
    increases our delta by the delta of the now-merged right-hand node.
    """
    if right is None:
      return self._left
    elif isinstance(right, PositionalIntersection) and right._left is None:
      self._delta += right._delta
      self._right = right._right
    else:
      self._right = right
    return self

  def _OpName(self):
    return "+%d" % self._delta

  def _SetOperation(self, callback, results):
    """Perform a standard intersection operation, but on a match,
    return only the left node's position data, and then only those
    positions with a difference of self._delta from the right node's
    positions.
    """
    assert len(results) == 2
    m1 = results[0].matches
    m2 = results[1].matches
    matches = []
    while m1 and m2:
      if m1[0] < m2[0]:
        m1 = m1[bisect_left(m1, m2[0]):]
      elif m2[0] < m1[0]:
        m2 = m2[bisect_left(m2, m1[0]):]
      else:
        new_data = [pos for pos in m1[0].data if (pos + self._delta in m2[0].data)]
        if new_data:
          matches.append(_MatchResult(key=m1[0].key, data=new_data))
        m1 = m1[1:]
        m2 = m2[1:]

    callback(_EvalResult(matches=matches, last_key=self._ComputeLastKey(results),
                         read_units=results[0].read_units + results[1].read_units))


class Parenthetical(_QueryNode):
  """This node encapsulates a child node and will be merged into
  _OpNodes as if it were a single value; This protects parenthesized
  trees from having their order adjusted.
  """
  def __init__(self, schema, child):
    # Collapse any nodes which have a missing left-hand node (this
    # can happen with a top-level positional intersection node that
    # has a place-holder term in the left-most position).
    if isinstance(child, PositionalIntersection) and child._left is None:
      child = child._right

    self._child = child
    self._precedence = 1000

  def PrintTree(self, level, param_dict):
    return self._child.PrintTree(level, param_dict)

  def Evaluate(self, client, callback, start_key, consistent_read, param_dict):
    self._child.Evaluate(client, callback, start_key, consistent_read, param_dict)


class PhraseNode(_QueryNode):
  """This node encapsulates a phrase, which may be a single term, or
  an arbitrarily complex subtree. For example, a phrase query with
  term expansions. The first pass through the parser merely sets the
  phrase on creation. However, when the tree is initialized, the
  phrase subtree is created.
  """
  def __init__(self, schema, table_name, column_name, phrase):
    """Convert the phrase into a query string using the appropriate
    indexer object (found via table:column in schema). This query
    string is passed to the phrase parser, and the new query subtree
    is set as the child of this PhraseNode.
    """
    self._precedence = 1000

    try:
      self.schema = schema
      self.table = self.schema.GetTable(table_name)
      assert isinstance(self.table, IndexedTable), table_name
      column = self.table.GetColumn(column_name)
      assert column.indexer, column_name
      self.column = column
      self.phrase = phrase
      self.child_parser = _PhraseParser(self.schema, self.table, column)
    except:
      logging.exception('phrase \'%s\'' % phrase)
      raise

  def _CreateChildNode(self, param_dict):
    if isinstance(self.phrase, Parameter):
      phrase = self.phrase.Resolve(param_dict)
    else:
      assert isinstance(self.phrase, basestring)
      phrase = self.phrase
    phrase_query_str = self.column.indexer.GetQueryString(self.column, phrase)
    match = _QUOTED_TOKEN_RE.match(phrase_query_str)
    if match is not None:
      # Fast path: it's one token, so create the IndexTermNode directly.
      return IndexTermNode(self.schema, self.table, self.column, match.group(1))
    else:
      # Slow path: run the real query parser to make sure we tokenize more complex queries correctly.
      return self.child_parser.Run(phrase_query_str)

  def PrintTree(self, level, param_dict):
    child = self._CreateChildNode(param_dict)
    return child.PrintTree(level, param_dict)