################################################################################
#  Licensed to the Apache Software Foundation (ASF) under one
#  or more contributor license agreements.  See the NOTICE file
#  distributed with this work for additional information
#  regarding copyright ownership.  The ASF licenses this file
#  to you 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 datetime
from enum import Enum
from functools import partial

from typing import Any, Tuple, Dict, List

from pyflink.common import Row
from pyflink.datastream.time_domain import TimeDomain
from pyflink.fn_execution import flink_fn_execution_pb2, pickle
from pyflink.serializers import PickleSerializer
from pyflink.table import functions
from pyflink.table.udf import DelegationTableFunction, DelegatingScalarFunction, \
    AggregateFunction, PandasAggregateFunctionWrapper

_func_num = 0
_constant_num = 0


def wrap_pandas_result(it):
    import pandas as pd
    return [pd.Series([result]) for result in it]


def check_pandas_udf_result(f, *input_args):
    output = f(*input_args)
    import pandas as pd
    assert type(output) == pd.Series or type(output) == pd.DataFrame, \
        "The result type of Pandas UDF '%s' must be pandas.Series or pandas.DataFrame, got %s" \
        % (f.__name__, type(output))
    assert len(output) == len(input_args[0]), \
        "The result length '%d' of Pandas UDF '%s' is not equal to the input length '%d'" \
        % (len(output), f.__name__, len(input_args[0]))
    return output


def extract_over_window_user_defined_function(user_defined_function_proto):
    window_index = user_defined_function_proto.window_index
    return (*extract_user_defined_function(user_defined_function_proto, True), window_index)


def extract_user_defined_function(user_defined_function_proto, pandas_udaf=False)\
        -> Tuple[str, Dict, List]:
    """
    Extracts user-defined-function from the proto representation of a
    :class:`UserDefinedFunction`.

    :param user_defined_function_proto: the proto representation of the Python
    :param pandas_udaf: whether the user_defined_function_proto is pandas udaf
    :class:`UserDefinedFunction`
    """

    def _next_func_num():
        global _func_num
        _func_num = _func_num + 1
        return _func_num

    def _extract_input(args) -> Tuple[str, Dict, List]:
        local_variable_dict = {}
        local_funcs = []
        args_str = []
        for arg in args:
            if arg.HasField("udf"):
                # for chaining Python UDF input: the input argument is a Python ScalarFunction
                udf_arg, udf_variable_dict, udf_funcs = extract_user_defined_function(arg.udf)
                args_str.append(udf_arg)
                local_variable_dict.update(udf_variable_dict)
                local_funcs.extend(udf_funcs)
            elif arg.HasField("inputOffset"):
                # the input argument is a column of the input row
                args_str.append("value[%s]" % arg.inputOffset)
            else:
                # the input argument is a constant value
                constant_value_name, parsed_constant_value = \
                    _parse_constant_value(arg.inputConstant)
                args_str.append(constant_value_name)
                local_variable_dict[constant_value_name] = parsed_constant_value
        return ",".join(args_str), local_variable_dict, local_funcs

    variable_dict = {}
    user_defined_funcs = []

    user_defined_func = pickle.loads(user_defined_function_proto.payload)
    if pandas_udaf:
        user_defined_func = PandasAggregateFunctionWrapper(user_defined_func)
    func_name = 'f%s' % _next_func_num()
    if isinstance(user_defined_func, DelegatingScalarFunction) \
            or isinstance(user_defined_func, DelegationTableFunction):
        if user_defined_function_proto.is_pandas_udf:
            variable_dict[func_name] = partial(check_pandas_udf_result, user_defined_func.func)
        else:
            variable_dict[func_name] = user_defined_func.func
    else:
        variable_dict[func_name] = user_defined_func.eval
    user_defined_funcs.append(user_defined_func)

    func_args, input_variable_dict, input_funcs = _extract_input(user_defined_function_proto.inputs)
    variable_dict.update(input_variable_dict)
    user_defined_funcs.extend(input_funcs)
    return "%s(%s)" % (func_name, func_args), variable_dict, user_defined_funcs


def _parse_constant_value(constant_value) -> Tuple[str, Any]:
    j_type = constant_value[0]
    serializer = PickleSerializer()
    pickled_data = serializer.loads(constant_value[1:])
    # the type set contains
    # TINYINT,SMALLINT,INTEGER,BIGINT,FLOAT,DOUBLE,DECIMAL,CHAR,VARCHAR,NULL,BOOLEAN
    # the pickled_data doesn't need to transfer to anther python object
    if j_type == 0:
        parsed_constant_value = pickled_data
    # the type is DATE
    elif j_type == 1:
        parsed_constant_value = \
            datetime.date(year=1970, month=1, day=1) + datetime.timedelta(days=pickled_data)
    # the type is TIME
    elif j_type == 2:
        seconds, milliseconds = divmod(pickled_data, 1000)
        minutes, seconds = divmod(seconds, 60)
        hours, minutes = divmod(minutes, 60)
        parsed_constant_value = datetime.time(hours, minutes, seconds, milliseconds * 1000)
    # the type is TIMESTAMP
    elif j_type == 3:
        parsed_constant_value = \
            datetime.datetime(year=1970, month=1, day=1, hour=0, minute=0, second=0) \
            + datetime.timedelta(milliseconds=pickled_data)
    else:
        raise Exception("Unknown type %s, should never happen" % str(j_type))

    def _next_constant_num():
        global _constant_num
        _constant_num = _constant_num + 1
        return _constant_num

    constant_value_name = 'c%s' % _next_constant_num()
    return constant_value_name, parsed_constant_value


def extract_user_defined_aggregate_function(
        current_index,
        user_defined_function_proto,
        distinct_info_dict: Dict[Tuple[List[str]], Tuple[List[int], List[int]]]):
    user_defined_agg = load_aggregate_function(user_defined_function_proto.payload)
    assert isinstance(user_defined_agg, AggregateFunction)
    args_str = []
    local_variable_dict = {}
    for arg in user_defined_function_proto.inputs:
        if arg.HasField("inputOffset"):
            # the input argument is a column of the input row
            args_str.append("value[%s]" % arg.inputOffset)
        else:
            # the input argument is a constant value
            constant_value_name, parsed_constant_value = \
                _parse_constant_value(arg.inputConstant)
            for key, value in local_variable_dict.items():
                if value == parsed_constant_value:
                    constant_value_name = key
                    break
            if constant_value_name not in local_variable_dict:
                local_variable_dict[constant_value_name] = parsed_constant_value
            args_str.append(constant_value_name)

    if user_defined_function_proto.distinct:
        if tuple(args_str) in distinct_info_dict:
            distinct_info_dict[tuple(args_str)][0].append(current_index)
            distinct_info_dict[tuple(args_str)][1].append(user_defined_function_proto.filter_arg)
            distinct_index = distinct_info_dict[tuple(args_str)][0][0]
        else:
            distinct_info_dict[tuple(args_str)] = \
                ([current_index], [user_defined_function_proto.filter_arg])
            distinct_index = current_index
    else:
        distinct_index = -1
    return user_defined_agg, \
        eval("lambda value : (%s,)" % ",".join(args_str), local_variable_dict) \
        if args_str else lambda v: tuple(), \
        user_defined_function_proto.filter_arg, \
        distinct_index


def is_built_in_function(payload):
    # The payload may be a pickled bytes or the class name of the built-in functions.
    # If it represents a built-in function, it will start with 0x00.
    # If it is a pickled bytes, it will start with 0x80.
    return payload[0] == 0


def load_aggregate_function(payload):
    if is_built_in_function(payload):
        built_in_function_class_name = payload[1:].decode("utf-8")
        cls = getattr(functions, built_in_function_class_name)
        return cls()
    else:
        return pickle.loads(payload)


def extract_data_stream_stateless_function(udf_proto):
    """
    Extracts user-defined-function from the proto representation of a
    :class:`Function`.

    :param udf_proto: the proto representation of the Python :class:`Function`
    """
    func_type = udf_proto.function_type
    UserDefinedDataStreamFunction = flink_fn_execution_pb2.UserDefinedDataStreamFunction
    func = None

    user_defined_func = pickle.loads(udf_proto.payload)
    if func_type == UserDefinedDataStreamFunction.MAP:
        func = user_defined_func.map
    elif func_type == UserDefinedDataStreamFunction.FLAT_MAP:
        func = user_defined_func.flat_map
    elif func_type == UserDefinedDataStreamFunction.REDUCE:
        reduce_func = user_defined_func.reduce

        def wrapped_func(value):
            return reduce_func(value[0], value[1])
        func = wrapped_func
    elif func_type == UserDefinedDataStreamFunction.CO_MAP:
        co_map_func = user_defined_func

        def wrapped_func(value):
            # value in format of: [INPUT_FLAG, REAL_VALUE]
            # INPUT_FLAG value of True for the left stream, while False for the right stream
            return Row(CoMapFunctionOutputFlag.LEFT.value, co_map_func.map1(value[1])) \
                if value[0] else Row(CoMapFunctionOutputFlag.RIGHT.value,
                                     co_map_func.map2(value[2]))
        func = wrapped_func
    elif func_type == UserDefinedDataStreamFunction.CO_FLAT_MAP:
        co_flat_map_func = user_defined_func

        def wrapped_func(value):
            if value[0]:
                result = co_flat_map_func.flat_map1(value[1])
                if result:
                    for result_val in result:
                        yield Row(CoFlatMapFunctionOutputFlag.LEFT.value, result_val)
                yield Row(CoFlatMapFunctionOutputFlag.LEFT_END.value, None)
            else:
                result = co_flat_map_func.flat_map2(value[2])
                if result:
                    for result_val in result:
                        yield Row(CoFlatMapFunctionOutputFlag.RIGHT.value, result_val)
                yield Row(CoFlatMapFunctionOutputFlag.RIGHT_END.value, None)
        func = wrapped_func

    elif func_type == UserDefinedDataStreamFunction.TIMESTAMP_ASSIGNER:
        extract_timestamp = user_defined_func.extract_timestamp

        def wrapped_func(value):
            pre_timestamp = value[0]
            real_data = value[1]
            return extract_timestamp(real_data, pre_timestamp)
        func = wrapped_func

    return func, user_defined_func


def extract_process_function(user_defined_function_proto, ctx):
    process_function = pickle.loads(user_defined_function_proto.payload)
    process_element = process_function.process_element

    def wrapped_process_function(value):
        # VALUE[CURRENT_TIMESTAMP, CURRENT_WATERMARK, NORMAL_DATA]
        ctx.set_timestamp(value[0])
        ctx.timer_service().set_current_watermark(value[1])
        output_result = process_element(value[2], ctx)
        return output_result

    return wrapped_process_function, process_function


def extract_keyed_process_function(user_defined_function_proto, ctx, on_timer_ctx,
                                   collector, keyed_state_backend):
    process_function = pickle.loads(user_defined_function_proto.payload)
    process_element = process_function.process_element
    on_timer = process_function.on_timer

    def wrapped_keyed_process_function(value):
        if value[0] is not None:
            # it is timer data
            # VALUE: TIMER_FLAG, TIMESTAMP_OF_TIMER, CURRENT_WATERMARK, CURRENT_KEY_OF_TIMER, None
            on_timer_ctx.set_timestamp(value[1])
            on_timer_ctx.timer_service().set_current_watermark(value[2])
            current_key = value[3]
            on_timer_ctx.set_current_key(current_key)
            keyed_state_backend.set_current_key(current_key)
            if value[0] == KeyedProcessFunctionInputFlag.EVENT_TIME_TIMER.value:
                on_timer_ctx.set_time_domain(TimeDomain.EVENT_TIME)
            elif value[0] == KeyedProcessFunctionInputFlag.PROC_TIME_TIMER.value:
                on_timer_ctx.set_time_domain(TimeDomain.PROCESSING_TIME)
            else:
                raise TypeError("TimeCharacteristic[%s] is not supported." % str(value[0]))
            output_result = on_timer(value[1], on_timer_ctx)
        else:
            # it is normal data
            # VALUE: TIMER_FLAG, CURRENT_TIMESTAMP, CURRENT_WATERMARK, None, NORMAL_DATA
            # NORMAL_DATA: CURRENT_KEY, DATA
            ctx.set_timestamp(value[1])
            ctx.timer_service().set_current_watermark(value[2])
            current_key = value[4][0]
            ctx.set_current_key(current_key)
            keyed_state_backend.set_current_key(Row(current_key))

            output_result = process_element(value[4][1], ctx)

        if output_result:
            for result in output_result:
                yield Row(None, None, None, result)

        for result in collector.buf:
            # 0: proc time timer data
            # 1: event time timer data
            # 2: normal data
            # result_row: [TIMER_FLAG, TIMER TYPE, TIMER_KEY, RESULT_DATA]
            yield Row(result[0], result[1], result[2], None)

        collector.clear()

    return wrapped_keyed_process_function, process_function


"""
All these Enum Classes MUST be in sync with
org.apache.flink.streaming.api.utils.PythonOperatorUtils if there are any changes.
"""


class KeyedProcessFunctionInputFlag(Enum):
    EVENT_TIME_TIMER = 0
    PROC_TIME_TIMER = 1
    NORMAL_DATA = 2


class KeyedProcessFunctionOutputFlag(Enum):
    REGISTER_EVENT_TIMER = 0
    REGISTER_PROC_TIMER = 1
    NORMAL_DATA = 2
    DEL_EVENT_TIMER = 3
    DEL_PROC_TIMER = 4


class CoFlatMapFunctionOutputFlag(Enum):
    LEFT = 0
    RIGHT = 1
    LEFT_END = 2
    RIGHT_END = 3


class CoMapFunctionOutputFlag(Enum):
    LEFT = 0
    RIGHT = 1