#!/usr/bin/env python3
# encoding: utf-8

import re
import itertools
import sys
import argparse

# Create SQL statement to verify dateTime64 is accepted as argument to functions taking DateTime.
FUNCTIONS = """
toTimeZone(N, 'UTC')
toYear(N, 'Asia/Istanbul')
toQuarter(N, 'Asia/Istanbul')
toMonth(N, 'Asia/Istanbul')
toDayOfYear(N, 'Asia/Istanbul')
toDayOfMonth(N, 'Asia/Istanbul')
toDayOfWeek(N, 0, 'Asia/Istanbul')
toHour(N, 'Asia/Istanbul')
toMinute(N, 'Asia/Istanbul')
toSecond(N, 'Asia/Istanbul')
toUnixTimestamp(N)
toStartOfYear(N, 'Asia/Istanbul')
toStartOfISOYear(N, 'Asia/Istanbul')
toStartOfQuarter(N, 'Asia/Istanbul')
toStartOfMonth(N, 'Asia/Istanbul')
toMonday(N, 'Asia/Istanbul')
toStartOfWeek(N, 'Asia/Istanbul')
toStartOfDay(N, 'Asia/Istanbul')
toStartOfHour(N, 'Asia/Istanbul')
toStartOfMinute(N, 'Asia/Istanbul')
toStartOfFiveMinutes(N, 'Asia/Istanbul')
toStartOfTenMinutes(N, 'Asia/Istanbul')
toStartOfFifteenMinutes(N, 'Asia/Istanbul')
toStartOfInterval(N, INTERVAL 1 year, 'Asia/Istanbul')
toStartOfInterval(N, INTERVAL 1 month, 'Asia/Istanbul')
toStartOfInterval(N, INTERVAL 1 day, 'Asia/Istanbul')
toStartOfInterval(N, INTERVAL 15 minute, 'Asia/Istanbul')
date_trunc('year', N, 'Asia/Istanbul')
date_trunc('month', N, 'Asia/Istanbul')
date_trunc('day', N, 'Asia/Istanbul')
date_trunc('minute', N, 'Asia/Istanbul')
toTime(N, 'Asia/Istanbul')
toRelativeYearNum(N, 'Asia/Istanbul')
toRelativeQuarterNum(N, 'Asia/Istanbul')
toRelativeMonthNum(N, 'Asia/Istanbul')
toRelativeWeekNum(N, 'Asia/Istanbul')
toRelativeDayNum(N, 'Asia/Istanbul')
toRelativeHourNum(N, 'Asia/Istanbul')
toRelativeMinuteNum(N, 'Asia/Istanbul')
toRelativeSecondNum(N, 'Asia/Istanbul')
toISOYear(N, 'Asia/Istanbul')
toISOWeek(N, 'Asia/Istanbul')
toWeek(N, 'Asia/Istanbul')
toYearWeek(N, 'Asia/Istanbul')
timeSlot(N, 'Asia/Istanbul')
toYYYYMM(N, 'Asia/Istanbul')
toYYYYMMDD(N, 'Asia/Istanbul')
toYYYYMMDDhhmmss(N, 'Asia/Istanbul')
addYears(N, 1, 'Asia/Istanbul')
addMonths(N, 1, 'Asia/Istanbul')
addWeeks(N, 1, 'Asia/Istanbul')
addDays(N, 1, 'Asia/Istanbul')
addHours(N, 1, 'Asia/Istanbul')
addMinutes(N, 1, 'Asia/Istanbul')
addSeconds(N, 1, 'Asia/Istanbul')
addQuarters(N, 1, 'Asia/Istanbul')
subtractYears(N, 1, 'Asia/Istanbul')
subtractMonths(N, 1, 'Asia/Istanbul')
subtractWeeks(N, 1, 'Asia/Istanbul')
subtractDays(N, 1, 'Asia/Istanbul')
subtractHours(N, 1, 'Asia/Istanbul')
subtractMinutes(N, 1, 'Asia/Istanbul')
subtractSeconds(N, 1, 'Asia/Istanbul')
subtractQuarters(N, 1, 'Asia/Istanbul')
CAST(N as DateTime('Europe/Minsk'))
CAST(N as Date)
CAST(N as UInt64)
CAST(N as DateTime64(0, 'Europe/Minsk'))
CAST(N as DateTime64(3, 'Europe/Minsk'))
CAST(N as DateTime64(6, 'Europe/Minsk'))
CAST(N as DateTime64(9, 'Europe/Minsk'))
# Casting our test values to DateTime(12) will cause an overflow and hence will fail the test under UB sanitizer.
# CAST(N as DateTime64(12, 'Asia/Istanbul'))
# DateTime64(18) will always fail due to zero precision, but it is Ok to test here:
# CAST(N as DateTime64(18, 'Asia/Istanbul'))
formatDateTime(N, '%C %d %D %e %F %H %I %j %m %M %p %R %S %T %u %V %w %y %Y %%', 'Asia/Istanbul')
""".splitlines()

# Expanded later to cartesian product of all arguments, using format string.
extra_ops = [
    # With same type:
    (
        ["N {op} N"],
        {
            "op": [
                "- ",  # does not work, but should it?
                "+ ",  # does not work, but should it?
                "!=",
                "==",  # equality and inequality supposed to take sub-second part in account
                "< ",
                "<=",
                "> ",
                ">=",
            ]
        },
    ),
    # With other DateTime types:
    (
        ["N {op} {arg}", "{arg} {op} N"],
        {
            "op": [
                "-",  # does not work, but should it?
                "!=",
                "==",
                # these are naturally expected to work, but they don't:
                "< ",
                "<=",
                "> ",
                ">=",
            ],
            "arg": ["DT", "D", "DT64"],
        },
    ),
    # With arithmetic types
    (
        ["N {op} {arg}", "{arg} {op} N"],
        {
            "op": ["+ ", "- ", "==", "!=", "< ", "<=", "> ", ">="],
            "arg": [
                "toUInt8(1)",
                "toInt8(-1)",
                "toUInt16(1)",
                "toInt16(-1)",
                "toUInt32(1)",
                "toInt32(-1)",
                "toUInt64(1)",
                "toInt64(-1)",
            ],
        },
    ),
]

# Expand extra_ops here
for funcs, args in extra_ops:
    args_keys = list(args.keys())
    for args_vals in itertools.product(*list(args.values())):
        for func in funcs:
            result_func = func.format(**dict(list(zip(args_keys, args_vals))))
            FUNCTIONS.append(result_func)

# filter out empty lines and commented out lines
COMMENTED_OUT_LINE_RE = re.compile(r"^\s*#")
FUNCTIONS = list(
    [f for f in FUNCTIONS if len(f) != 0 and COMMENTED_OUT_LINE_RE.match(f) == None]
)
TYPES = ["D", "DT", "DT64"]


def escape_string(s):
    if sys.version_info[0] > 2:
        return s.encode("unicode_escape").decode("utf-8").replace("'", "\\'")
    else:
        return s.encode("string-escape").decode("utf-8")


def execute_functions_for_types(functions, types):
    # NOTE: use string.Template here to allow lines with missing keys, like type, e.g. SELECT CAST(toDateTime64(1234567890), 'DateTime64')
    for func in functions:
        print(("""SELECT 'SELECT {func}';""".format(func=escape_string(func))))
        for dt in types:
            prologue = "\
WITH \
toDateTime64('2019-09-16 19:20:11.234', 3, 'Europe/Minsk') as DT64, \
toDateTime('2019-09-16 19:20:11', 'Europe/Minsk') as DT, \
toDate('2019-09-16') as D, {X} as N".format(
                X=dt
            )
            print(
                (
                    """{prologue} SELECT toTypeName(r), {func} as r FORMAT CSV;""".format(
                        prologue=prologue, func=func
                    )
                )
            )
        print("""SELECT '------------------------------------------';""")


def main():
    def parse_args():
        parser = argparse.ArgumentParser()
        parser.add_argument(
            "--functions_re",
            type=re.compile,
            help="RE to enable functions",
            default=None,
        )
        parser.add_argument(
            "--types_re",
            type=lambda s: re.compile("^(" + s + ")$"),
            help="RE to enable types, supported types: " + ",".join(TYPES),
            default=None,
        )
        parser.add_argument(
            "--list_functions",
            action="store_true",
            help="List all functions to be tested and exit",
        )
        return parser.parse_args()

    args = parse_args()

    functions = FUNCTIONS
    types = TYPES

    if args.functions_re:
        functions = list([f for f in functions if args.functions_re.search(f)])
        if len(functions) == 0:
            print("functions list is empty")
            return -1

    if args.types_re:
        types = list([t for t in types if args.types_re.match(t)])
        if len(types) == 0:
            print("types list is empty")
            return -1

    if args.list_functions:
        print(("\n".join(functions)))
        return 0

    execute_functions_for_types(functions, types)


if __name__ == "__main__":
    exit(main())