# Copyright 2014 Cloudera Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import re
import locale
import string
import platform
import warnings
import operator

from functools import reduce

import sqlalchemy as sa
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql.functions import GenericFunction

from ibis.sql.alchemy import (
    unary, varargs, fixed_arity, Over, _variance_reduction, _get_sqla_table
)
import ibis.expr.analytics as L
import ibis.expr.datatypes as dt
import ibis.expr.operations as ops
import ibis.expr.window as W

import ibis.sql.alchemy as alch
_operation_registry = alch._operation_registry.copy()


# TODO: substr and find are copied from SQLite, we should really have a
# "base" set of SQL functions that are the most common APIs across the major
# RDBMS
def _substr(t, expr):
    f = sa.func.substr

    arg, start, length = expr.op().args

    sa_arg = t.translate(arg)
    sa_start = t.translate(start)

    if length is None:
        return f(sa_arg, sa_start + 1)
    else:
        sa_length = t.translate(length)
        return f(sa_arg, sa_start + 1, sa_length)


def _string_find(t, expr):
    arg, substr, start, _ = expr.op().args

    if start is not None:
        raise NotImplementedError

    sa_arg = t.translate(arg)
    sa_substr = t.translate(substr)

    return sa.func.strpos(sa_arg, sa_substr) - 1


def _infix_op(infix_sym):
    def formatter(t, expr):
        op = expr.op()
        left, right = op.args

        left_arg = t.translate(left)
        right_arg = t.translate(right)
        return left_arg.op(infix_sym)(right_arg)

    return formatter


def _extract(fmt):
    def translator(t, expr):
        arg, = expr.op().args
        sa_arg = t.translate(arg)
        return sa.extract(fmt, sa_arg)
    return translator


def _cast(t, expr):
    arg, typ = expr.op().args

    sa_arg = t.translate(arg)
    sa_type = t.get_sqla_type(typ)

    # specialize going from an integer type to a timestamp
    if isinstance(arg.type(), dt.Integer) and isinstance(sa_type, sa.DateTime):
        return sa.func.timezone('UTC', sa.func.to_timestamp(sa_arg))

    if arg.type().equals(dt.binary) and typ.equals(dt.string):
        return sa.func.encode(sa_arg, 'escape')

    if typ.equals(dt.binary):
        #  decode yields a column of memoryview which is annoying to deal with
        # in pandas. CAST(expr AS BYTEA) is correct and returns byte strings.
        return sa.cast(sa_arg, sa.Binary())

    return sa.cast(sa_arg, sa_type)


def _typeof(t, expr):
    arg, = expr.op().args
    sa_arg = t.translate(arg)
    typ = sa.cast(sa.func.pg_typeof(sa_arg), sa.TEXT)

    # select pg_typeof('asdf') returns unknown so we have to check the child's
    # type for nullness
    return sa.case(
        [
            ((typ == 'unknown') & (arg.type() != dt.null), 'text'),
            ((typ == 'unknown') & (arg.type() == dt.null), 'null'),
        ],
        else_=typ
    )


def _second(t, expr):
    # extracting the second gives us the fractional part as well, so smash that
    # with a cast to SMALLINT
    sa_arg, = map(t.translate, expr.op().args)
    return sa.cast(sa.extract('second', sa_arg), sa.SMALLINT)


def _millisecond(t, expr):
    # we get total number of milliseconds including seconds with extract so we
    # mod 1000
    sa_arg, = map(t.translate, expr.op().args)
    return sa.cast(sa.extract('millisecond', sa_arg), sa.SMALLINT) % 1000


def _string_agg(t, expr):
    # we could use sa.func.string_agg since postgres 9.0, but we can cheaply
    # maintain backwards compatibility here, so we don't use it
    arg, sep = map(t.translate, expr.op().args)
    return sa.func.array_to_string(sa.func.array_agg(arg), sep)


_strftime_to_postgresql_rules = {
    '%a': 'TMDy',  # TM does it in a locale dependent way
    '%A': 'TMDay',
    '%w': 'D',  # 1-based day of week, see below for how we make this 0-based
    '%d': 'DD',  # day of month
    '%-d': 'FMDD',  # - is no leading zero for Python same for FM in postgres
    '%b': 'TMMon',  # Sep
    '%B': 'TMMonth',  # September
    '%m': 'MM',  # 01
    '%-m': 'FMMM',  # 1
    '%y': 'YY',  # 15
    '%Y': 'YYYY',  # 2015
    '%H': 'HH24',  # 09
    '%-H': 'FMHH24',  # 9
    '%I': 'HH12',  # 09
    '%-I': 'FMHH12',  # 9
    '%p': 'AM',  # AM or PM
    '%M': 'MI',  # zero padded minute
    '%-M': 'FMMI',  # Minute
    '%S': 'SS',  # zero padded second
    '%-S': 'FMSS',  # Second
    '%f': 'US',  # zero padded microsecond
    '%z': 'OF',  # utf offset
    '%Z': 'TZ',  # uppercase timezone name
    '%j': 'DDD',  # zero padded day of year
    '%-j': 'FMDDD',  # day of year
    '%U': 'WW',  # 1-based week of year
    # 'W': ?,  # meh
}

try:
    _strftime_to_postgresql_rules.update({
        '%c': locale.nl_langinfo(locale.D_T_FMT),  # locale date and time
        '%x': locale.nl_langinfo(locale.D_FMT),  # locale date
        '%X': locale.nl_langinfo(locale.T_FMT)  # locale time
    })
except AttributeError:
    warnings.warn(
        'locale specific date formats (%%c, %%x, %%X) are not yet implemented '
        'for %s' % platform.system()
    )


def tokenize_noop(scanner, token):
    return token


# translate strftime spec into mostly equivalent PostgreSQL spec
_scanner = re.Scanner([
    (py, tokenize_noop)
    for py in _strftime_to_postgresql_rules.keys()
] + [
    # "%e" is in the C standard and Python actually generates this if your spec
    # contains "%c" but we don't officially support it as a specifier so we
    # need to special case it in the scanner
    (r'%e', tokenize_noop),

    # double quotes need to be escaped
    (r'"', lambda scanner, token: re.escape(token)),

    # spaces should be greedily consumed and kept
    (r'\s+', tokenize_noop),

    (r'[%s]' % re.escape(string.punctuation), tokenize_noop),

    # everything else except double quotes and spaces
    (r'[^%s\s]+' % re.escape(string.punctuation), tokenize_noop),
])


_lexicon_values = frozenset(_strftime_to_postgresql_rules.values())

_strftime_blacklist = frozenset(['%w', '%U', '%c', '%x', '%X', '%e'])


def _reduce_tokens(tokens, arg):
    # current list of tokens
    curtokens = []

    # reduced list of tokens that accounts for blacklisted values
    reduced = []

    non_special_tokens = (
        frozenset(_strftime_to_postgresql_rules) - _strftime_blacklist
    )

    # TODO: how much of a hack is this?
    for token in tokens:

        # we are a non-special token %A, %d, etc.
        if token in non_special_tokens:
            curtokens.append(_strftime_to_postgresql_rules[token])

        # we have a string like DD, to escape this we
        # surround it with double quotes
        elif token in _lexicon_values:
            curtokens.append('"%s"' % token)

        # we have a token that needs special treatment
        elif token in _strftime_blacklist:
            if token == '%w':
                value = sa.extract('dow', arg)  # 0 based day of week
            elif token == '%U':
                value = sa.cast(sa.func.to_char(arg, 'WW'), sa.SMALLINT) - 1
            elif token == '%c' or token == '%x' or token == '%X':
                # re scan and tokenize this pattern
                try:
                    new_pattern = _strftime_to_postgresql_rules[token]
                except KeyError:
                    raise ValueError(
                        'locale specific date formats (%%c, %%x, %%X) are '
                        'not yet implemented for %s' % platform.system()
                    )

                new_tokens, _ = _scanner.scan(new_pattern)
                value = reduce(
                    sa.sql.ColumnElement.concat,
                    _reduce_tokens(new_tokens, arg)
                )
            elif token == '%e':
                # pad with spaces instead of zeros
                value = sa.func.replace(sa.func.to_char(arg, 'DD'), '0', ' ')

            reduced += [
                sa.func.to_char(arg, ''.join(curtokens)),
                sa.cast(value, sa.TEXT)
            ]

            # empty current token list in case there are more tokens
            del curtokens[:]

        # uninteresting text
        else:
            curtokens.append(token)
    else:
        # append result to r if we had more tokens or if we have no
        # blacklisted tokens
        if curtokens:
            reduced.append(sa.func.to_char(arg, ''.join(curtokens)))
    return reduced


def _strftime(t, expr):
    arg, pattern = map(t.translate, expr.op().args)
    tokens, _ = _scanner.scan(pattern.value)
    reduced = _reduce_tokens(tokens, arg)
    return reduce(sa.sql.ColumnElement.concat, reduced)


def _distinct_from(a, b):
    return (((a != None) | (b != None)) &  # noqa
            sa.func.coalesce(a != b, True))


def _find_in_set(t, expr):
    # postgresql 9.5 has array_position, but the code below works on any
    # version of postgres with generate_subscripts
    # TODO: could make it even more generic by using generate_series
    # TODO: this works with *any* type, not just strings. should the operation
    #       itself also have this property?
    arg, haystack = expr.op().args
    needle = t.translate(arg)
    haystack = sa.select([sa.literal(
        [element._arg.value for element in haystack],
        type_=sa.dialects.postgresql.ARRAY(needle.type)
    ).label('haystack')]).cte()

    subscripts = sa.select([
        sa.func.generate_subscripts(haystack.c.haystack, 1).label('i')
    ]).cte()

    # return a zero based index
    return sa.select([subscripts.c.i - 1]).where(
        ~_distinct_from(haystack.c.haystack[subscripts.c.i], needle)
    ).order_by(subscripts.c.i).limit(1)


def _regex_replace(t, expr):
    string, pattern, replacement = map(t.translate, expr.op().args)

    # postgres defaults to replacing only the first occurrence
    return sa.func.regexp_replace(string, pattern, replacement, 'g')


def _reduction(func_name):
    def reduction_compiler(t, expr):
        arg, where = expr.op().args

        if arg.type().equals(dt.boolean):
            arg = arg.cast('int32')

        func = getattr(sa.func, func_name)

        if where is not None:
            arg = where.ifelse(arg, None)
        return func(t.translate(arg))
    return reduction_compiler


def _log(t, expr):
    arg, base = expr.op().args
    arg = t.translate(arg)
    if base is not None:
        return sa.func.log(t.translate(base), arg)
    else:
        return sa.func.ln(arg)


_cumulative_to_reduction = {
    ops.CumulativeSum: ops.Sum,
    ops.CumulativeMin: ops.Min,
    ops.CumulativeMax: ops.Max,
    ops.CumulativeMean: ops.Mean,
    ops.CumulativeAny: ops.Any,
    ops.CumulativeAll: ops.All,
}


def _cumulative_to_window(translator, expr, window):
    win = W.cumulative_window()
    win = win.group_by(window._group_by).order_by(window._order_by)

    op = expr.op()

    klass = _cumulative_to_reduction[type(op)]
    new_op = klass(*op.args)
    new_expr = expr._factory(new_op, name=expr._name)

    if type(new_op) in translator._rewrites:
        new_expr = translator._rewrites[type(new_op)](new_expr)

    return L.windowize_function(new_expr, win)


def _window(t, expr):
    op = expr.op()

    arg, window = op.args
    reduction = t.translate(arg)

    window_op = arg.op()

    _require_order_by = (
        ops.Lag,
        ops.Lead,
        ops.DenseRank,
        ops.MinRank,
        ops.NTile,
        ops.PercentRank,
        ops.FirstValue,
        ops.LastValue,
    )

    if isinstance(window_op, ops.CumulativeOp):
        arg = _cumulative_to_window(t, arg, window)
        return t.translate(arg)

    # Some analytic functions need to have the expression of interest in
    # the ORDER BY part of the window clause
    if isinstance(window_op, _require_order_by) and not window._order_by:
        order_by = t.translate(window_op.args[0])
    else:
        order_by = list(map(t.translate, window._order_by))

    partition_by = list(map(t.translate, window._group_by))

    result = Over(
        reduction,
        partition_by=partition_by,
        order_by=order_by,
        preceding=window.preceding,
        following=window.following,
    )

    if isinstance(
        window_op, (ops.RowNumber, ops.DenseRank, ops.MinRank, ops.NTile)
    ):
        return result - 1
    else:
        return result


def _ntile(t, expr):
    op = expr.op()
    args = op.args
    arg, buckets = map(t.translate, args)
    return sa.func.ntile(buckets)


class regex_extract(GenericFunction):
    def __init__(self, string, pattern, index):
        super(regex_extract, self).__init__(string, pattern, index)
        self.string = string
        self.pattern = pattern
        self.index = index


@compiles(regex_extract, 'postgresql')
def compile_regex_extract(element, compiler, **kw):
    return '(REGEXP_MATCHES(%s, %s))[%s]' % (
        compiler.process(element.string, **kw),
        compiler.process(element.pattern, **kw),
        compiler.process(element.index, **kw),
    )


def _regex_extract(t, expr):
    string, pattern, index = map(t.translate, expr.op().args)
    return sa.case(
        [
            (
                sa.func.textregexeq(string, pattern),
                sa.func.regex_extract(string, pattern, index + 1)
            )
        ],
        else_=''
    )


def _cardinality(array):
    return sa.case(
        [(array == None, None)],  # noqa: E711
        else_=sa.func.coalesce(sa.func.array_length(array, 1), 0),
    )


def _array_repeat(t, expr):
    """Is this really that useful?

    Repeat an array like a Python list using modular arithmetic,
    scalar subqueries, and PostgreSQL's ARRAY function.

    This is inefficient if PostgreSQL allocates memory for the entire sequence
    and the output column. A quick glance at PostgreSQL's C code shows the
    sequence is evaluated stepwise, which suggests that it's roughly constant
    memory for the sequence generation.
    """
    raw, times = map(t.translate, expr.op().args)

    # SQLAlchemy uses our column's table in the FROM clause. We need a simpler
    # expression to workaround this.
    array = sa.column(raw.name, type_=raw.type)

    # We still need to prefix the table name to the column name in the final
    # query, so make sure the column knows its origin
    array.table = raw.table

    array_length = _cardinality(array)

    # sequence from 1 to the total number of elements desired in steps of 1.
    # the call to greatest isn't necessary, but it provides clearer intent
    # rather than depending on the implicit postgres generate_series behavior
    start = step = 1
    stop = sa.func.greatest(times, 0) * array_length
    series = sa.func.generate_series(start, stop, step).alias()
    series_column = sa.column(series.name, type_=sa.INTEGER)

    # if our current index modulo the array's length
    # is a multiple of the array's length, then the index is the array's length
    index_expression = series_column % array_length
    index = sa.func.coalesce(sa.func.nullif(index_expression, 0), array_length)

    # tie it all together in a scalar subquery and collapse that into an ARRAY
    selected = sa.select([array[index]]).select_from(series)
    return sa.func.array(selected.as_scalar())


def _identical_to(t, expr):
    left, right = args = expr.op().args
    if left.equals(right):
        return True
    else:
        left, right = map(t.translate, args)
        return left.op('IS NOT DISTINCT FROM')(right)


def _hll_cardinality(t, expr):
    # postgres doesn't have a builtin HLL algorithm, so we default to standard
    # count distinct for now
    arg, _ = expr.op().args
    sa_arg = t.translate(arg)
    return sa.func.count(sa.distinct(sa_arg))


def _table_column(t, expr):
    op = expr.op()
    ctx = t.context
    table = op.table

    sa_table = _get_sqla_table(ctx, table)
    out_expr = getattr(sa_table.c, op.name)

    expr_type = expr.type()

    if isinstance(expr_type, dt.Timestamp):
        timezone = expr_type.timezone
        if timezone is not None:
            out_expr = out_expr.op('AT TIME ZONE')(timezone).label(op.name)

    # If the column does not originate from the table set in the current SELECT
    # context, we should format as a subquery
    if t.permit_subquery and ctx.is_foreign_expr(table):
        return sa.select([out_expr])

    return out_expr


def _round(t, expr):
    arg, digits = expr.op().args
    sa_arg = t.translate(arg)

    if digits is None:
        return sa.func.round(sa_arg)

    # postgres doesn't allow rounding of double precision values to a specific
    # number of digits (though simple truncation on doubles is allowed) so
    # we cast to numeric and then cast back if necessary
    result = sa.func.round(sa.cast(sa_arg, sa.NUMERIC), t.translate(digits))
    if digits is not None and isinstance(arg.type(), dt.Decimal):
        return result
    return sa.cast(result, sa.dialects.postgresql.DOUBLE_PRECISION())


def _mod(t, expr):
    left, right = map(t.translate, expr.op().args)

    # postgres doesn't allow modulus of double precision values, so upcast and
    # then downcast later if necessary
    if not isinstance(expr.type(), dt.Integer):
        left = sa.cast(left, sa.NUMERIC)
        right = sa.cast(right, sa.NUMERIC)

    result = left % right
    if expr.type().equals(dt.double):
        return sa.cast(result, sa.dialects.postgresql.DOUBLE_PRECISION())
    else:
        return result


def _floor_divide(t, expr):
    left, right = map(t.translate, expr.op().args)
    return sa.func.floor(left / right)


def _array_slice(t, expr):
    arg, start, stop = expr.op().args
    sa_arg = t.translate(arg)
    sa_start = t.translate(start)

    if stop is None:
        sa_stop = _cardinality(sa_arg)
    else:
        sa_stop = t.translate(stop)
    return sa_arg[sa_start + 1:sa_stop]


_operation_registry.update({
    # We override this here to support time zones
    ops.TableColumn: _table_column,

    # types
    ops.Cast: _cast,
    ops.TypeOf: _typeof,

    # miscellaneous varargs
    ops.Least: varargs(sa.func.least),
    ops.Greatest: varargs(sa.func.greatest),

    # null handling
    ops.IfNull: fixed_arity(sa.func.coalesce, 2),

    # boolean reductions
    ops.Any: fixed_arity(sa.func.bool_or, 1),
    ops.All: fixed_arity(sa.func.bool_and, 1),
    ops.NotAny: fixed_arity(lambda x: sa.not_(sa.func.bool_or(x)), 1),
    ops.NotAll: fixed_arity(lambda x: sa.not_(sa.func.bool_and(x)), 1),

    # strings
    ops.Substring: _substr,
    ops.StrRight: fixed_arity(sa.func.right, 2),
    ops.StringFind: _string_find,
    ops.StringLength: unary('length'),
    ops.GroupConcat: _string_agg,
    ops.Lowercase: unary('lower'),
    ops.Uppercase: unary('upper'),
    ops.Strip: unary('trim'),
    ops.LStrip: unary('ltrim'),
    ops.RStrip: unary('rtrim'),
    ops.LPad: fixed_arity('lpad', 3),
    ops.RPad: fixed_arity('rpad', 3),
    ops.Reverse: unary('reverse'),
    ops.Capitalize: unary('initcap'),
    ops.Repeat: fixed_arity('repeat', 2),
    ops.StringReplace: fixed_arity(sa.func.replace, 3),
    ops.StringSQLLike: _infix_op('LIKE'),
    ops.RegexSearch: _infix_op('~'),
    ops.RegexReplace: _regex_replace,
    ops.Translate: fixed_arity('translate', 3),
    ops.StringAscii: fixed_arity(sa.func.ascii, 1),
    ops.RegexExtract: _regex_extract,

    ops.FindInSet: _find_in_set,

    ops.Ceil: fixed_arity(sa.func.ceil, 1),
    ops.Floor: fixed_arity(sa.func.floor, 1),
    ops.FloorDivide: _floor_divide,
    ops.Exp: fixed_arity(sa.func.exp, 1),
    ops.Sign: fixed_arity(sa.func.sign, 1),
    ops.Sqrt: fixed_arity(sa.func.sqrt, 1),
    ops.Log: _log,
    ops.Ln: fixed_arity(sa.func.ln, 1),
    ops.Log2: fixed_arity(lambda x: sa.func.log(2, x), 1),
    ops.Log10: fixed_arity(sa.func.log, 1),
    ops.Power: fixed_arity(sa.func.power, 2),
    ops.Round: _round,
    ops.Modulus: _mod,

    # dates and times
    ops.Strftime: _strftime,
    ops.ExtractYear: _extract('year'),
    ops.ExtractMonth: _extract('month'),
    ops.ExtractDay: _extract('day'),
    ops.ExtractHour: _extract('hour'),
    ops.ExtractMinute: _extract('minute'),
    ops.ExtractSecond: _second,
    ops.ExtractMillisecond: _millisecond,
    ops.Sum: _reduction('sum'),
    ops.Mean: _reduction('avg'),
    ops.Min: _reduction('min'),
    ops.Max: _reduction('max'),
    ops.Variance: _variance_reduction('var'),
    ops.StandardDev: _variance_reduction('stddev'),

    # now is in the timezone of the server, but we want UTC
    ops.TimestampNow: lambda *args: sa.func.timezone('UTC', sa.func.now()),
    ops.WindowOp: _window,
    ops.CumulativeOp: _window,
    ops.RowNumber: fixed_arity(lambda: sa.func.row_number(), 0),
    ops.DenseRank: fixed_arity(lambda arg: sa.func.dense_rank(), 1),
    ops.MinRank: fixed_arity(lambda arg: sa.func.rank(), 1),
    ops.PercentRank: fixed_arity(lambda arg: sa.func.percent_rank(), 1),
    ops.NTile: _ntile,

    # array operations
    ops.ArrayLength: fixed_arity(_cardinality, 1),
    ops.ArrayCollect: fixed_arity(sa.func.array_agg, 1),
    ops.ArraySlice: _array_slice,
    ops.ArrayIndex: fixed_arity(lambda array, index: array[index + 1], 2),
    ops.ArrayConcat: fixed_arity(operator.add, 2),
    ops.ArrayRepeat: _array_repeat,
    ops.IdenticalTo: _identical_to,
    ops.HLLCardinality: _hll_cardinality,
})


def add_operation(op, translation_func):
    _operation_registry[op] = translation_func


class PostgreSQLExprTranslator(alch.AlchemyExprTranslator):

    _registry = _operation_registry
    _rewrites = alch.AlchemyExprTranslator._rewrites.copy()
    _type_map = alch.AlchemyExprTranslator._type_map.copy()
    _type_map.update({
        dt.Double: sa.types.FLOAT,
        dt.Float: sa.types.REAL
    })


rewrites = PostgreSQLExprTranslator.rewrites
compiles = PostgreSQLExprTranslator.compiles


class PostgreSQLDialect(alch.AlchemyDialect):

    translator = PostgreSQLExprTranslator
