#!/usr/bin/env python
# vim: ts=4 et sw=4 sts=4 syntax=python

"""jsonsql

Copyright (c) 2015, Los Alamos National Security, LLC
All rights reserved.

Copyright (2015). Los Alamos National Security, LLC. This software was produced
under U.S. Government contract DE-AC52-06NA25396 for Los Alamos National
Laboratory (LANL), which is operated by Los Alamos National Security, LLC for
the U.S. Department of Energy. The U.S. Government has rights to use,
reproduce, and distribute this software. NEITHER THE GOVERNMENT NOR LOS ALAMOS
NATIONAL SECURITY, LLC MAKES ANY WARRANTY, EXPRESS OR IMPLIED, OR ASSUMES ANY
LIABILITY FOR THE USE OF THIS SOFTWARE. If software is modified to produce
derivative works, such modified software should be clearly marked, so as not to
confuse it with the version available from LANL.

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
of the Software, and to permit persons to whom the Software is furnished to do
so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

SQL query evaluation on flat files of JSON records using JavaScript.

Author: Curt Hash <chash@lanl.gov>

"""

import argparse
import os
import subprocess
import sys
import tempfile

from pyparsing import alphanums, alphas, CaselessKeyword, Combine, \
    delimitedList, Forward, Keyword, nums, oneOf, opAssoc, \
    operatorPrecedence, Optional, Or, ParseException, ParserElement, \
    QuotedString, Regex, Suppress, Word, ZeroOrMore


# This enables memoization in pyparsing. It's very slow without it.
ParserElement.enablePackrat()


# JS code declaring global variables and functions.
BOILERPLATE = r'''
// This program was generated by jsonsql.

var ip = require('ip');
var Netmask = require('netmask').Netmask;
var strftime = require('strftime');

var records = 0;
var partitions = {};
var state = {};
var distinct = {};
var stop = false;

// Throws an error if `a` is not an array.
function assert_array(a) {
  if (!Array.isArray(a)) {
    throw "assert_array() failed: " + JSON.stringify(a);
  }
}

// Returns a RegExp object equivalent to the LIKE expression `expr`.
function like2regex(expr) {
  expr = expr.replace(/[-\/\\^$*+?.()|[\]{}]/g, '\\$&');

  var head;
  if (expr[0] != '%') {
    head = '^';
  } else {
    head = '';
    expr = substr(expr, 2);
  }

  var tail;
  if (expr[-1] != '%') {
    tail = '$';
  } else {
    tail = '';
    expr = substr(expr, 1, expr.length - 1);
  }

  expr = expr.replace(/%/g, '.*');

  expr = head + expr + tail;

  return new RegExp(expr);
}

// Returns the concatenation of the arrays within the array.
function flatten(a) {
  assert_array(a);
  assert_array(a[0]);

  for (var i=1; i<a.length; i++) {
    assert_array(a[i]);
    a[0] = a[0].concat(a[i]);
  }

  return a[0];
}

// Returns the concatenation of the array arguments.
function concat() {
  assert_array(arguments[0]);

  for (var i=1; i<arguments.length; i++) {
    assert_array(arguments[i]);
    arguments[0] = arguments[0].concat(arguments[i]);
  }

  return arguments[0];
}

// Returns the number `n` truncated to the nearest integer.
function int(n) {
  return n | 0;
}

// Returns the epoch timestamp `epoch` formatted according to the `format`
// string.
function strftime_wrapper(format, epoch) {
  return strftime(format, new Date(epoch * 1000));
}

// Returns true if `addr` is in the CIDR block, false otherwise.
function cidr_contains(cidr, addr) {
  var block = new Netmask(cidr);
  return block.contains(addr);
}

// Returns the substring of `s` matched by the `n`-th (starting at 1)
// parenthesized group of the regular expression `r`.
function submatch(s, r, n) {
  return s.match(r)[n];
}

// Returns the sum of the array values.
function array_sum(a) {
  assert_array(a);

  var sum = 0;
  for (var i=0; i<a.length; i++) {
    sum += a[i];
  }

  return sum;
}

// Returns the average of the array values.
function array_avg(a) {
  assert_array(a);

  if (a.length == 0) {
    throw "array_avg() is not defined on an empty array";
  }

  return array_sum(a) / a.length;
}

// Returns the maximum of the array values.
function array_max(a) {
  assert_array(a);

  if (a.length == 0) {
    throw "array_max() is not defined on an empty array";
  }

  var max = Number.NEGATIVE_INFINITY;
  for (var i=0; i<a.length; i++) {
    max = Math.max(max, a[i]);
  }

  return max;
}

// Returns the minimum of the array values.
function array_min(a) {
  assert_array(a);

  if (a.length == 0) {
    throw "array_min() is not defined on an empty array";
  }

  var min = Number.POSITIVE_INFINITY;
  for (var i=0; i<a.length; i++) {
    min = Math.min(min, a[i]);
  }

  return min;
}

// Returns true if any values in the array `a` evaluate to true, false
// otherwise.
function any(a) {
  assert_array(a);

  for (var i=0; i<a.length; i++) {
    if (a[i]) {
      return true;
    }
  }

  return false;
}

// Returns true if all values in the array `a` evaluate to true, false
// otherwise.
function all(a) {
  assert_array(a);

  for (var i=0; i<a.length; i++) {
    if (!a[i]) {
      return false;
    }
  }

  return true;
}
'''


# Code for functions defined at runtime.
DYNAMIC_FUNCTIONS = []


class Identifier(object):

    """Identifier. """

    def __init__(self, tokens):
        self.identifier = tokens[0]
        self.aggregate = False

    def js_code(self):
        """Returns the identifier. """
        return self.identifier


class NumericLiteral(object):

    """Numeric literal. """

    def __init__(self, tokens):
        self.number = tokens[0]
        self.aggregate = False

    def js_code(self):
        """Returns the number. """
        return self.number


class StringLiteral(object):

    """String literal. """

    def __init__(self, tokens):
        self.string = tokens[0]
        self.aggregate = False

    def js_code(self):
        """Returns the quoted string. """
        return "'%s'" % self.string


class RegexLiteral(object):

    """Regex literal. """

    def __init__(self, tokens):
        self.regex = tokens[0]
        self.aggregate = False

    def js_code(self):
        """Returns the pattern. """
        return '/%s/' % self.regex


class FunctionLiteral(object):

    """Reference to a function. """

    def __init__(self, tokens):
        self.function = tokens[0]
        self.aggregate = False

    def js_code(self):
        """Returns the function name. """
        return self.function


# Maps JSONPath expressions to variable names.
JSON_PATHS = {}


class JSONPath(object):

    """JSONPath expression. """

    def __init__(self, tokens):
        self.expression = tokens[0]
        self.aggregate = False

        # Figure out a key name for the expression in case it's projected.
        self.name = self.expression.split('.')[-1].split('[')[0]
        if self.name == '*':
            self.name = 'items'

        # Map the expression to a variable name. Variables are used so that
        # each expression is evaluated only once for each input record.
        json_path = "jsonpath.eval(obj, '%s')" % self.expression
        if not json_path in JSON_PATHS:
            variable = 'jp%d' % len(JSON_PATHS)
            JSON_PATHS[json_path] = variable
        self.variable = JSON_PATHS[json_path]

    def js_code(self):
        """Returns a reference to the evaluated expression. """
        return self.variable


class UnaryExpression(object):

    """Unary expression. """

    def __init__(self, tokens):
        self.oper, self.operand = tokens[0]
        self.aggregate = self.operand.aggregate

    def js_code(self):
        """Returns the equivalent JS expression. """
        assert not self.aggregate

        return self.oper + self.operand.js_code()

    def js_update(self):
        """Returns JS code that updates the incremental state of the aggregate
        operand.

        """
        assert self.aggregate

        return self.operand.js_update()

    def js_final(self):
        """Returns JS code that evaluates the final value of the aggregate
        expression.

        """
        assert self.aggregate

        return self.oper + self.operand.js_final()


class BinaryExpression(object):

    """Binary expression. """

    def __init__(self, tokens):
        self.tokens = tokens[0]
        self.operand = self.tokens[0]

        # Maps SQL operators to their JS equivalents.
        operator_map = {
            '=' : '==',
            '||' : '+',
            'and' : '&&',
            'or' : '||'
        }

        rest = self.tokens[1:]
        self.pairs = [(operator_map.get(oper.lower(), oper.lower()), operand)
                      for oper, operand in zip(rest[::2], rest[1::2])]

        self.aggregate = (self.operand.aggregate or
                          any(operand.aggregate for _, operand in self.pairs))

    def js_code(self):
        """Returns the equivalent JS expression. """
        assert not self.aggregate

        expr = [self.operand.js_code()]
        for oper, operand in self.pairs:
            value = operand.js_code()

            if oper.endswith('like') or oper[-1] == '~':
                # Translate the LIKE/match expression into a RegExp test.
                lvalue = expr.pop()

                if oper.endswith('like'):
                    # LIKE or NOT LIKE
                    test = 'like2regex(%s).test(%s)' % (value, lvalue)
                else:
                    # ~ or !~
                    test = '%s.test(%s)' % (value, lvalue)

                if oper in ['not like', '!~']:
                    test = '!' + test

                expr.append(test)
            else:
                expr.append(oper)
                expr.append(value)

        return '(' + ' '.join(expr) + ')'

    def js_update(self):
        """Returns JS code that updates the incremental state of the aggregate
        operand(s).

        """
        assert self.aggregate

        code = []

        if self.operand.aggregate:
            code.append(self.operand.js_update())

        for _, operand in self.pairs:
            if operand.aggregate:
                code.append(operand.js_update())

        return '\n'.join(code)

    def js_final(self):
        """Returns JS code that evaluates the final value of the aggregate
        expression.

        """
        assert self.aggregate

        if self.operand.aggregate:
            operand = self.operand.js_final()
        else:
            operand = self.operand.js_code()

        expr = [operand]
        for oper, operand in self.pairs:
            if operand.aggregate:
                value = operand.js_final()
            else:
                value = operand.js_code()

            if oper.endswith('like') or oper[-1] == '~':
                # Translate the LIKE/match expression into a RegExp test.
                lvalue = expr.pop()

                if oper.endswith('like'):
                    # LIKE or NOT LIKE
                    test = 'like2regex(%s).test(%s)' % (value, lvalue)
                else:
                    # ~ or !~
                    test = '%s.test(%s)' % (value, lvalue)

                if oper in ['not like', '!~']:
                    test = '!' + test

                expr.append(test)
            else:
                expr.append(oper)
                expr.append(value)

        return '(%s)' % ' '.join(expr)


# Each aggregate function uses a key to access its state information. This list
# is used to initialize the state object for each function.
STATE_KEYS = []


class AggregateFunction(object):

    """Aggregate function (e.g., count). """

    def __init__(self, tokens):
        self.function = tokens[0].lower()

        if type(tokens[1]) == str and tokens[1].lower() == 'distinct':
            self.distinct = True
            self.arg = tokens[2]
        else:
            self.distinct = False
            self.arg = tokens[1]

        self._statekey = "'agg%d'" % len(STATE_KEYS)
        STATE_KEYS.append(self._statekey)

        self.aggregate = True

    def js_update(self):
        """Returns JS code that updates the incremental state of the aggregate
        function.

        The code should reference the current partition using the `part`
        variable.

        """
        if self.arg == '*':
            value = 'line'
        else:
            value = self.arg.js_code()

        code = ['''
            if (!(part in state[%(SK)s])) {
              // Initialize state for the new partition.
              state[%(SK)s][part] = {
                distinct: {},
                count: 0,
                sum: 0,
                max: Number.NEGATIVE_INFINITY,
                min: Number.POSITIVE_INFINITY
              };
            }
            ''' % {'SK' : self._statekey}]

        if self.distinct:
            code.append('''
                // Distinct check for %(FUNC)s().'
                if (state[%(SK)s][part]['distinct'][%(VALUE)s] === undefined) {
                  state[%(SK)s][part]['distinct'][%(VALUE)s] = 1;
                ''' % {
                    'FUNC' : self.function,
                    'SK' : self._statekey,
                    'VALUE' : value
                })

        if self.function == 'count':
            code.append('''
                // Update incremental state for count().
                state[%(SK)s][part]['count']++;
                ''' % {'SK' : self._statekey})
        elif self.function == 'avg':
            code.append('''
                // Update incremental state for avg().
                state[%(SK)s][part]['count']++;
                state[%(SK)s][part]['sum'] += %(VALUE)s;
                ''' % {
                    'SK' : self._statekey,
                    'VALUE' : value
                })
        elif self.function == 'max':
            code.append('''
                // Update incremental state for max().
                if (%(VALUE)s > state[%(SK)s][part]['max']) {
                  state[%(SK)s][part]['max'] = %(VALUE)s;
                }
                ''' % {
                    'SK' : self._statekey,
                    'VALUE' : value
                })
        elif self.function == 'min':
            code.append('''
                // Update incremental state for min().
                if (%(VALUE)s < state[%(SK)s][part]['min']) {
                  state[%(SK)s][part]['min'] = %(VALUE)s;
                }
                ''' % {
                    'SK' : self._statekey,
                    'VALUE' : value
                })
        elif self.function in ['sum', 'total']:
            code.append('''
                // Update incremental state for sum().
                state[%(SK)s][part]['sum'] += %(VALUE)s;
                ''' % {
                    'SK' : self._statekey,
                    'VALUE' : value
                })
        else:
            raise NotImplementedError(self.function)

        if self.distinct:
            code.append('} // End distinct check for %s().' % self.function)

        return '\n'.join(code)

    def js_final(self):
        """Returns JS code that evaluates the final value of the aggregate
        function.

        This code should reference the current partition using the `part`
        variable.

        """
        if self.function == 'count':
            # count() should return 0 if no rows were processed.
            return ("(!(part in state[%(SK)s])) ? "
                    "0 : state[%(SK)s][part]['count']"
                    % {'SK' : self._statekey})

        code = ("(!(part in state[%s])) ? (function() {throw '%s() is not "
                "defined on an empty partition; partition=' + part}) : "
                % (self._statekey, self.function))

        if self.function == 'avg':
            code += ("state[%(SK)s][part]['sum'] / "
                     "state[%(SK)s][part]['count']" % {'SK' : self._statekey})
        elif self.function == 'max':
            code += "state[%s][part]['max']" % self._statekey
        elif self.function == 'min':
            code += "state[%s][part]['min']" % self._statekey
        elif self.function in ['sum', 'total']:
            code += "state[%s][part]['sum']" % self._statekey
        else:
            raise NotImplementedError(self.function)

        return code


class Function(object):

    """Non-aggregate function. """

    def __init__(self, tokens):
        self.function = tokens[0].lower()
        self.args = tokens[1:]
        self.aggregate = any(arg.aggregate for arg in self.args)

        self._function_map = {
            'abs' : 'Math.abs',
            'max' : 'Math.max',
            'min' : 'Math.min',
            'strftime' : 'strftime_wrapper',
            'ip_mask' : 'ip.mask',
            'lower' : 'toLowerCase',
            'upper' : 'toUpperCase'
        }

        self._array_functions = ['array_avg', 'array_len', 'array_max',
                                 'array_min', 'array_sum', 'concat', 'flatten',
                                 'map']

        self._method_functions = ['ltrim', 'replace', 'rtrim', 'substr',
                                  'toLowerCase', 'toUpperCase', 'trim']

    def _call_code(self, function, args):
        """Returns function call code. """
        if function in ['length', 'array_len']:
            # Translate to length property of first argument.
            return '%s.length' % args[0]

        if function in self._method_functions:
            # Translate to method of first argument.
            return '%s.%s(%s)' % (args[0], function, ', '.join(args[1:]))

        if function == 'map':
            # Define the map function.
            function = 'map_df%d' % len(DYNAMIC_FUNCTIONS)

            map_function = args.pop(0)
            map_function = self._function_map.get(map_function, map_function)

            map_code = '''
               function %s(a) {
                 assert_array(a);

                 var r = new Array();
                 for (var i=0; i<a.length; i++) {
                   r[i] = %s;
                 }

                 return r;
               }
               ''' % (function, self._call_code(map_function, ['a[i]']))

            DYNAMIC_FUNCTIONS.append(map_code)

        return '%s(%s)' % (function, ', '.join(args))

    def js_code(self):
        """Returns the equivalent JS expression. """
        assert not self.aggregate

        function = self._function_map.get(self.function, self.function)
        args = [a.js_code() for a in self.args]

        return self._call_code(function, args)

    def js_update(self):
        """Returns JS code that updates the incremental state of the aggregate
        argument(s).

        """
        assert self.aggregate

        return '\n'.join(a.js_update() for a in self.args if a.aggregate)

    def js_final(self):
        """Returns JS code that evaluates the final value of the function. """
        assert self.aggregate

        function = self._function_map.get(self.function, self.function)

        args = []
        for arg in self.args:
            if arg.aggregate:
                args.append(arg.js_final())
            else:
                args.append(arg.js_code())

        return self._call_code(function, args)


class Projection(object):

    """Field expression. """

    names = set()

    def __init__(self, tokens):
        self.expression = tokens[0]
        if self.expression == '*':
            # '*' is a special projection that is handled later.
            return

        self.alias = None
        if len(tokens) > 1:
            self.name = self.alias = tokens[2].identifier
        elif type(self.expression) == JSONPath:
            self.name = self.expression.name
        elif type(self.expression) in [AggregateFunction, Function]:
            self.name = self.expression.function
        else:
            self.name = 'expr'

        # Don't allow two projections to have the same key name.
        if self.name in self.__class__.names:
            raise ValueError('duplicate output key name: "%s"' % self.name)
        self.__class__.names.add(self.name)

        self.aggregate = self.expression.aggregate

    def js_code(self):
        """Returns the equivalent JS expression. """
        assert not self.aggregate

        return self.expression.js_code()

    def js_update(self):
        """Returns JS code that updates the incremental state of the aggregate
        expression.

        """
        assert self.aggregate

        return self.expression.js_update()

    def js_final(self):
        """Returns JS code that evaluates the final value of the aggregate
        expression.

        """
        assert self.aggregate

        return self.expression.js_final()


class ArrayComprehension(object):

    """Array comprehension. """

    def __init__(self, tokens):
        self.expr = tokens[0]
        self.identifier = tokens[1]
        self.array = tokens[2]

        self.aggregate = self.expr.aggregate or self.array.aggregate

        if len(tokens) == 4:
            self.filter_expr = tokens[3]
            self.aggregate = self.aggregate or self.filter_expr.aggregate
        else:
            self.filter_expr = None

    def js_code(self):
        """Returns JS code that evaluates the list comprehension. """
        if self.filter_expr:
            filter_expr = self.filter_expr.js_code()
        else:
            filter_expr = 'true'

        function = 'df%d_lc' % len(DYNAMIC_FUNCTIONS)
        function_code = '''
            function %s(__a) {
              assert_array(__a);

              var __r = new Array();
              for (var __i=0; __i<__a.length; __i++) {
                var %s = __a[__i];
                if (%s) {
                  __r.push(%s);
                }
              }

              return __r;
            }
            ''' % (function, self.identifier.js_code(), filter_expr,
                   self.expr.js_code())
        DYNAMIC_FUNCTIONS.append(function_code)

        return '%s(%s)' % (function, self.array.js_code())

    def js_update(self):
        """Returns JS code that updates the incremental state of the aggregate
        expression(s).

        """
        code = []

        if self.expr.aggregate:
            code.append(self.expr.js_update())

        if self.array.aggregate:
            code.append(self.expr.js_update())

        if self.filter_expr and self.filter_expr.aggregate:
            code.append(self.filter_expr.js_update())

        return '\n'.join(code)

    def js_final(self):
        """Returns JS code that evaluates the final value of the list
        comprehension.

        """
        if self.filter_expr:
            if self.filter_expr.aggregate:
                filter_expr = self.filter_expr.js_final()
            else:
                filter_expr = self.filter_expr.js_code()
        else:
            filter_expr = 'true'

        if self.expr.aggregate:
            expr = self.expr.js_final()
        else:
            expr = self.expr.js_code()

        if self.array.aggregate:
            array = self.array.js_final()
        else:
            array = self.array.js_code()

        function = 'df%d_lc' % len(DYNAMIC_FUNCTIONS)
        function_code = '''
            function %s(__a) {
              assert_array(__a);

              var __r = new Array();
              for (var __i=0; __i<__a.length; __i++) {
                var %s = __a[__i];
                if (%s) {
                  __r.push(%s);
                }
              }

              return __r;
            }
            ''' % (function, self.identifier.identifier, filter_expr, expr)
        DYNAMIC_FUNCTIONS.append(function_code)

        return '%s(%s)' % (function, array)


class Where(object):

    """Filter expression. """

    def __init__(self, tokens):
        self.expression = tokens[0]
        self.aggregate = False

    def js_code(self):
        """Returns the equivalent JS expression. """
        return self.expression.js_code()


def grammar():
    """Returns the SQL query statement grammar. """
    lpar = Suppress('(')
    rpar = Suppress(')')
    lbrack = Suppress('[')
    rbrack = Suppress(']')
    star = Keyword('*')

    def split_keywords(string):
        """Returns a dictionary mapping the tokens in `string` to
        CaselessKeyword objects.

        """
        return dict(((w, CaselessKeyword(w)) for w in string.split()))

    functions = ('abs all any array_avg array_len array_max array_min '
                 'array_sum cidr_contains concat flatten int ip_mask length '
                 'lower ltrim map max min replace rtrim strftime submatch '
                 'substr trim upper')

    functions = split_keywords(functions)
    function_name = Or(functions.values())
    function_literal = function_name.copy().setParseAction(FunctionLiteral)

    aggregates = split_keywords('avg count max min sum total')
    aggregate_function_name = Or(aggregates.values())

    keywords = split_keywords('in where as distinct and or like select limit '
                              'for if')
    keywords['not like'] = CaselessKeyword('not like')
    keyword = Or(keywords.values())

    json_path = Regex(r'\$(\.{1,2}'
                      r'(\w+(\[.+?\])?|'
                      r'\*))+').setParseAction(JSONPath)

    numeric_literal = Combine(
        (Word(nums) + Optional('.' + Word(nums)) | '.' + Word(nums)) +
        Optional('E' + Optional(oneOf('+ -')) + Word(nums))
    ).setParseAction(NumericLiteral)
    string_literal = (QuotedString("'") |
                      QuotedString('"')).setParseAction(StringLiteral)
    regex_literal = QuotedString('/').setParseAction(RegexLiteral)
    literal_value = (numeric_literal | string_literal | regex_literal |
                     function_literal)

    expr = Forward()

    identifier = (~keyword +
                  Word(alphas, alphanums + '_')).setParseAction(Identifier)
    array_comp = (lbrack + expr + Suppress(keywords['for']) + identifier +
                  Suppress(keywords['in']) + expr +
                  Optional(Suppress(keywords['if']) + expr) +
                  rbrack).setParseAction(ArrayComprehension)

    aggregate_function_arg = Optional(keywords['distinct']) + expr
    count_function_arg = Optional(keywords['distinct']) + (expr | star)
    aggregate_function = (aggregate_function_name + lpar +
                          aggregate_function_arg + rpar | aggregates['count'] +
                          lpar + count_function_arg +
                          rpar).setParseAction(AggregateFunction)

    function_arg = Optional(delimitedList(expr))
    function = (function_name + lpar + Optional(function_arg) +
                rpar).setParseAction(Function)

    expr_term = (array_comp | aggregate_function | function | literal_value |
                 json_path | identifier | lpar + expr + rpar)
    expr << operatorPrecedence(expr_term, [
        ('^', 2, opAssoc.LEFT, BinaryExpression),
        (oneOf('- + !'), 1, opAssoc.RIGHT, UnaryExpression),
        ('||', 2, opAssoc.LEFT, BinaryExpression),
        (oneOf('* / %'), 2, opAssoc.LEFT, BinaryExpression),
        (oneOf('+ -'), 2, opAssoc.LEFT, BinaryExpression),
        (oneOf('< <= > >='), 2, opAssoc.LEFT, BinaryExpression),
        (oneOf('= == !='), 2, opAssoc.LEFT, BinaryExpression),
        (oneOf('~ !~'), 2, opAssoc.LEFT, BinaryExpression),
        (keywords['not like'], 2, opAssoc.LEFT, BinaryExpression),
        (keywords['like'], 2, opAssoc.LEFT, BinaryExpression),
        (keywords['and'], 2, opAssoc.LEFT, BinaryExpression),
        (keywords['or'], 2, opAssoc.LEFT, BinaryExpression),
    ])

    all_fields = star.copy()
    projection = (all_fields | (expr + Optional(keywords['as'] +
                                identifier))).setParseAction(Projection)

    where_expr = (keywords['where'] +
                  expr.setResultsName('where').setParseAction(Where))

    limit_expr = (keywords['limit'] + Word(nums).setResultsName('limit'))

    projection_list = projection + ZeroOrMore(Suppress(',') + projection)

    query_stmt = Optional(keywords['select']) + \
                 Optional(keywords['distinct']).setResultsName('distinct') + \
                 projection_list.setResultsName('projections') + \
                 Optional(where_expr) + Optional(limit_expr)

    return query_stmt


class Query(object):

    """Parsed query text.

    Attributes:

        where : a Where object or None
        projections : a list of Projection objects or None if '*'
        star : True if * is being projected
        aggregate : True if an aggregate function is used, False otherwise
        distinct: True if distinct is used, False otherwise
        limit: integer row limit or None

    """

    def __init__(self, query):
        query_stmt = grammar()

        query = query.strip()
        if not query or query[:5].lower() in ['where', 'limit']:
            # Implicitly project all columns.
            query = ('* ' + query).strip()

        try:
            parsed = query_stmt.parseString(query, parseAll=True).asDict()
        except ParseException as err:
            print
            print query
            print ' ' * err.loc + '^'
            print err.msg
            raise

        self.distinct = 'distinct' in parsed

        projections = parsed['projections']
        self.projections = [p for p in projections if p.expression != '*']
        self.star = len(self.projections) != len(projections)

        self.aggregate = any(p.aggregate for p in self.projections)

        if 'where' in parsed:
            self.where = parsed['where']
        else:
            self.where = None

        if 'limit' in parsed:
            self.limit = int(parsed['limit'])
            if self.limit <= 0:
                raise ValueError('limit must be a positive integer')
        else:
            self.limit = None

    def js_code(self):
        """Returns the equivalent JS program as a string. """
        # Declare global variables and functions.
        code = [BOILERPLATE]

        # Initialize state variables for each aggregate.
        for key in STATE_KEYS:
            code.append('state[%s] = {};' % key)

        # Declare JSONPath variables. These are declared outside of the line
        # handler function so that their final values are accessible to the
        # close handler function (e.g., for evaluating count(*) + <field>).
        code.extend('var %s;' % v for v in JSON_PATHS.values())

        # Line processing code.
        code.append('''
            var jsonpath = require('JSONPath');

            var readline = require('readline').createInterface({
              input: process.stdin,
              output: process.stdout,
              terminal: false
            }).on('line', function(line) { // Line handler.
            ''')

        if self.limit:
            code.append('''
                if (stop) {
                  // The close() event doesn't happen immediately. Prevent
                  // processing additional lines.
                  return;
                }
                ''')

        code.append('var obj = JSON.parse(line);')

        # Evaluate JSONPath expressions once and reference with variables.
        code.extend('%s = %s;' % (v, e) for e, v in JSON_PATHS.items())

        if self.where:
            # Filter records with an if statement.
            code.append('if (%s) { // Filter.' % self.where.js_code())

        if not self.aggregate:
            # No aggregate projections.

            # Evaluate and output all projections.
            if self.star:
                code.append('var record = obj;')
            else:
                code.append('var record = {};')

            code.extend("record['%s'] = %s;" % (p.name, p.js_code())
                        for p in self.projections)

            code.append('''
                json = JSON.stringify(record);
                if (json != '{}') {
                ''')

            if self.distinct:
                # Only output a record if it hasn't been seen before.
                code.append('''
                    if (!(json in distinct)) { // Distinct check.
                        distinct[json] = 1;
                    ''')

            code.append('console.log(json);')

            if self.limit:
                code.append('''
                    records++;
                    if (records == %s) {
                      // Limit reached.
                      this.close();
                      stop = 1;
                    }
                    ''' % self.limit)

            if self.distinct:
                code.append('} // End distinct check.')

            code.append('}')
        else:
            # Evaluate aggregate projections.

            aggregates = []
            non_aggregates = []
            for projection in self.projections:
                if projection.aggregate:
                    aggregates.append(projection)
                else:
                    non_aggregates.append(projection)

            if non_aggregates:
                # Partition on non-aggregate projections.
                if self.star:
                    code.append('var part = obj;')
                else:
                    code.append('var part = {};')

                code.extend("part['%s'] = %s;" % (p.name, p.js_code())
                            for p in non_aggregates)
            else:
                code.append('var part = {};')

            # Set the partition and add it to the array of partitions.
            code.append('''
                part = JSON.stringify(part);
                partitions[part] = 1;
                ''')

            # Incrementally evaluate the aggregate projections.
            for projection in aggregates:
                code.append(projection.js_update())

        if self.where:
            # Close the filter if statement.
            code.append('} // End filter.')

        code.append('}); // End line handler.')

        if self.aggregate:
            # Output a record for each partition.
            code.append('''
                readline.on('close', function() {
                  if (Object.keys(partitions).length == 0) {
                    // No rows passed the filter, but some aggregate functions
                    // are valid on 0 rows.
                    partitions['{}'] = 1;
                  }

                  for (part in partitions) {
                    var record = JSON.parse(part);
                ''')

            code.extend("record['%s'] = %s;" % (p.name, p.js_final())
                        for p in self.projections if p.aggregate)

            code.append('console.log(JSON.stringify(record));')

            if self.limit:
                code.append('''
                    records++;
                    if (records == %s) {
                      // Limit reached.
                      return;
                    }
                    ''' % self.limit)

            code.append('''
                  } // End partitions loop.
                }); // End close handler.
                ''')

        # Define dynamic functions.
        code.extend(DYNAMIC_FUNCTIONS)

        return '\n'.join(code)


def main():
    """Parses args and executes. """
    parser = argparse.ArgumentParser()
    parser.add_argument('query', nargs=1, help='SQL query')
    parser.add_argument('-D', '--debug', action='store_true', default=False,
                        help='write the JS program to stdout and exit')
    args = parser.parse_args()

    query = Query(args.query[0])

    code = query.js_code()
    if args.debug:
        print code
        return

    # Write the JS code to a temporary file.
    with tempfile.NamedTemporaryFile(prefix='jsonsql.') as tmpjs:
        tmpjs.write(code)
        tmpjs.flush()

        # Expect node modules to live in the same directory as this script.
        env = os.environ.copy()
        node_dir = os.path.dirname(os.path.realpath(__file__))
        env['NODE_PATH'] = os.path.join(node_dir, 'node_modules')

        # Execute the JS code.
        node = subprocess.Popen(['nodejs', tmpjs.name], stdin=sys.stdin,
                                stdout=sys.stdout, stderr=sys.stderr, env=env)
        status = node.wait()

    return status


if __name__ == '__main__':
    sys.exit(main())
