################################################################################
#  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 abc
from functools import reduce
from itertools import chain
from typing import Tuple

from pyflink.fn_execution.coders import DataViewFilterCoder, PickleCoder
from pyflink.fn_execution.datastream.timerservice import InternalTimer
from pyflink.fn_execution.datastream.operations import Operation
from pyflink.fn_execution.datastream.timerservice_impl import TimerOperandType, InternalTimerImpl
from pyflink.fn_execution.table.state_data_view import extract_data_view_specs

from pyflink.fn_execution.table.window_assigner import TumblingWindowAssigner, \
    CountTumblingWindowAssigner, SlidingWindowAssigner, CountSlidingWindowAssigner, \
    SessionWindowAssigner
from pyflink.fn_execution.table.window_trigger import EventTimeTrigger, ProcessingTimeTrigger, \
    CountTrigger
from pyflink.fn_execution.utils import operation_utils
from pyflink.fn_execution.utils.operation_utils import extract_user_defined_aggregate_function

try:
    from pyflink.fn_execution.table.aggregate_fast import RowKeySelector, \
        SimpleAggsHandleFunction, GroupAggFunction, DistinctViewDescriptor, \
        SimpleTableAggsHandleFunction, GroupTableAggFunction
    from pyflink.fn_execution.table.window_aggregate_fast import \
        SimpleNamespaceAggsHandleFunction, GroupWindowAggFunction
    from pyflink.fn_execution.coder_impl_fast import InternalRow
    has_cython = True
except ImportError:
    from pyflink.fn_execution.table.aggregate_slow import RowKeySelector, \
        SimpleAggsHandleFunction, GroupAggFunction, DistinctViewDescriptor, \
        SimpleTableAggsHandleFunction, GroupTableAggFunction
    from pyflink.fn_execution.table.window_aggregate_slow import \
        SimpleNamespaceAggsHandleFunction, GroupWindowAggFunction
    has_cython = False

from pyflink.table import FunctionContext, Row

# UDF
SCALAR_FUNCTION_URN = "flink:transform:scalar_function:v1"

# UDTF
TABLE_FUNCTION_URN = "flink:transform:table_function:v1"

# UDAF
STREAM_GROUP_AGGREGATE_URN = "flink:transform:stream_group_aggregate:v1"
STREAM_GROUP_TABLE_AGGREGATE_URN = "flink:transform:stream_group_table_aggregate:v1"
STREAM_GROUP_WINDOW_AGGREGATE_URN = "flink:transform:stream_group_window_aggregate:v1"

# Pandas UDAF
PANDAS_AGGREGATE_FUNCTION_URN = "flink:transform:aggregate_function:arrow:v1"
PANDAS_BATCH_OVER_WINDOW_AGGREGATE_FUNCTION_URN = \
    "flink:transform:batch_over_window_aggregate_function:arrow:v1"


class BundleOperation(object):
    def finish_bundle(self):
        raise NotImplementedError


class BaseOperation(Operation):
    def __init__(self, serialized_fn):
        super(BaseOperation, self).__init__(serialized_fn)
        self.func, self.user_defined_funcs = self.generate_func(serialized_fn)

    def process_element(self, value):
        return self.func(value)

    def open(self):
        for user_defined_func in self.user_defined_funcs:
            if hasattr(user_defined_func, 'open'):
                user_defined_func.open(FunctionContext(self.base_metric_group))

    def close(self):
        for user_defined_func in self.user_defined_funcs:
            if hasattr(user_defined_func, 'close'):
                user_defined_func.close()

    @abc.abstractmethod
    def generate_func(self, serialized_fn) -> Tuple:
        pass


class ScalarFunctionOperation(BaseOperation):
    def __init__(self, serialized_fn, one_arg_optimization=False, one_result_optimization=False):
        self._one_arg_optimization = one_arg_optimization
        self._one_result_optimization = one_result_optimization
        super(ScalarFunctionOperation, self).__init__(serialized_fn)

    def generate_func(self, serialized_fn):
        """
        Generates a lambda function based on udfs.
        :param serialized_fn: serialized function which contains a list of the proto
                              representation of the Python :class:`ScalarFunction`
        :return: the generated lambda function
        """
        scalar_functions, variable_dict, user_defined_funcs = reduce(
            lambda x, y: (
                ','.join([x[0], y[0]]),
                dict(chain(x[1].items(), y[1].items())),
                x[2] + y[2]),
            [operation_utils.extract_user_defined_function(
                udf, one_arg_optimization=self._one_arg_optimization)
                for udf in serialized_fn.udfs])
        if self._one_result_optimization:
            func_str = 'lambda value: %s' % scalar_functions
        else:
            func_str = 'lambda value: [%s]' % scalar_functions
        generate_func = eval(func_str, variable_dict)
        return generate_func, user_defined_funcs


class TableFunctionOperation(BaseOperation):
    def __init__(self, serialized_fn):
        super(TableFunctionOperation, self).__init__(serialized_fn)

    def generate_func(self, serialized_fn):
        """
        Generates a lambda function based on udtfs.
        :param serialized_fn: serialized function which contains the proto representation of
                              the Python :class:`TableFunction`
        :return: the generated lambda function
        """
        table_function, variable_dict, user_defined_funcs = \
            operation_utils.extract_user_defined_function(serialized_fn.udfs[0])
        variable_dict['normalize_table_function_result'] = \
            operation_utils.normalize_table_function_result
        generate_func = eval('lambda value: normalize_table_function_result(%s)' % table_function,
                             variable_dict)
        return generate_func, user_defined_funcs


class PandasAggregateFunctionOperation(BaseOperation):
    def __init__(self, serialized_fn):
        super(PandasAggregateFunctionOperation, self).__init__(serialized_fn)

    def generate_func(self, serialized_fn):
        pandas_functions, variable_dict, user_defined_funcs = reduce(
            lambda x, y: (
                ','.join([x[0], y[0]]),
                dict(chain(x[1].items(), y[1].items())),
                x[2] + y[2]),
            [operation_utils.extract_user_defined_function(udf, True)
             for udf in serialized_fn.udfs])
        variable_dict['normalize_pandas_result'] = operation_utils.normalize_pandas_result
        generate_func = eval('lambda value: normalize_pandas_result([%s])' %
                             pandas_functions, variable_dict)
        return generate_func, user_defined_funcs


class PandasBatchOverWindowAggregateFunctionOperation(BaseOperation):
    def __init__(self, serialized_fn):
        super(PandasBatchOverWindowAggregateFunctionOperation, self).__init__(serialized_fn)
        self.windows = [window for window in serialized_fn.windows]
        # the index among all the bounded range over window
        self.bounded_range_window_index = [-1 for _ in range(len(self.windows))]
        # Whether the specified position window is a bounded range window.
        self.is_bounded_range_window = []

        from pyflink.fn_execution import flink_fn_execution_pb2

        window_types = flink_fn_execution_pb2.OverWindow

        bounded_range_window_nums = 0
        for i, window in enumerate(self.windows):
            window_type = window.window_type
            if (window_type is window_types.RANGE_UNBOUNDED_PRECEDING) or (
                    window_type is window_types.RANGE_UNBOUNDED_FOLLOWING) or (
                    window_type is window_types.RANGE_SLIDING):
                self.bounded_range_window_index[i] = bounded_range_window_nums
                self.is_bounded_range_window.append(True)
                bounded_range_window_nums += 1
            else:
                self.is_bounded_range_window.append(False)

    def generate_func(self, serialized_fn):
        user_defined_funcs = []
        self.window_indexes = []
        self.mapper = []
        for udf in serialized_fn.udfs:
            pandas_agg_function, variable_dict, user_defined_func, window_index = \
                operation_utils.extract_over_window_user_defined_function(udf)
            user_defined_funcs.extend(user_defined_func)
            self.window_indexes.append(window_index)
            self.mapper.append(eval('lambda value: %s' % pandas_agg_function, variable_dict))
        return self.wrapped_over_window_function, user_defined_funcs

    def wrapped_over_window_function(self, boundaries_series):
        import pandas as pd
        from pyflink.fn_execution import flink_fn_execution_pb2

        OverWindow = flink_fn_execution_pb2.OverWindow
        input_series = boundaries_series[-1]
        # the row number of the arrow format data
        input_cnt = len(input_series[0])
        results = []
        # loop every agg func
        for i in range(len(self.window_indexes)):
            window_index = self.window_indexes[i]
            # the over window which the agg function belongs to
            window = self.windows[window_index]
            window_type = window.window_type
            func = self.mapper[i]
            result = []
            if self.is_bounded_range_window[window_index]:
                window_boundaries = boundaries_series[
                    self.bounded_range_window_index[window_index]]
                if window_type is OverWindow.RANGE_UNBOUNDED_PRECEDING:
                    # range unbounded preceding window
                    for j in range(input_cnt):
                        end = window_boundaries[j]
                        series_slices = [s.iloc[:end] for s in input_series]
                        result.append(func(series_slices))
                elif window_type is OverWindow.RANGE_UNBOUNDED_FOLLOWING:
                    # range unbounded following window
                    for j in range(input_cnt):
                        start = window_boundaries[j]
                        series_slices = [s.iloc[start:] for s in input_series]
                        result.append(func(series_slices))
                else:
                    # range sliding window
                    for j in range(input_cnt):
                        start = window_boundaries[j * 2]
                        end = window_boundaries[j * 2 + 1]
                        series_slices = [s.iloc[start:end] for s in input_series]
                        result.append(func(series_slices))
            else:
                # unbounded range window or unbounded row window
                if (window_type is OverWindow.RANGE_UNBOUNDED) or (
                        window_type is OverWindow.ROW_UNBOUNDED):
                    series_slices = [s.iloc[:] for s in input_series]
                    func_result = func(series_slices)
                    result = [func_result for _ in range(input_cnt)]
                elif window_type is OverWindow.ROW_UNBOUNDED_PRECEDING:
                    # row unbounded preceding window
                    window_end = window.upper_boundary
                    for j in range(input_cnt):
                        end = min(j + window_end + 1, input_cnt)
                        series_slices = [s.iloc[: end] for s in input_series]
                        result.append(func(series_slices))
                elif window_type is OverWindow.ROW_UNBOUNDED_FOLLOWING:
                    # row unbounded following window
                    window_start = window.lower_boundary
                    for j in range(input_cnt):
                        start = max(j + window_start, 0)
                        series_slices = [s.iloc[start: input_cnt] for s in input_series]
                        result.append(func(series_slices))
                else:
                    # row sliding window
                    window_start = window.lower_boundary
                    window_end = window.upper_boundary
                    for j in range(input_cnt):
                        start = max(j + window_start, 0)
                        end = min(j + window_end + 1, input_cnt)
                        series_slices = [s.iloc[start: end] for s in input_series]
                        result.append(func(series_slices))
            results.append(pd.Series(result))
        return results


class BaseStatefulOperation(BaseOperation, abc.ABC):

    def __init__(self, serialized_fn, keyed_state_backend):
        self.keyed_state_backend = keyed_state_backend
        super(BaseStatefulOperation, self).__init__(serialized_fn)

    def finish(self):
        super().finish()
        if self.keyed_state_backend:
            self.keyed_state_backend.commit()


NORMAL_RECORD = 0
TRIGGER_TIMER = 1
REGISTER_EVENT_TIMER = 0
REGISTER_PROCESSING_TIMER = 1


class AbstractStreamGroupAggregateOperation(BaseStatefulOperation):

    def __init__(self, serialized_fn, keyed_state_backend):
        self.generate_update_before = serialized_fn.generate_update_before
        self.grouping = [i for i in serialized_fn.grouping]
        self.group_agg_function = None
        # If the upstream generates retract message, we need to add an additional count1() agg
        # to track current accumulated messages count. If all the messages are retracted, we need
        # to send a DELETE message to downstream.
        self.index_of_count_star = serialized_fn.index_of_count_star
        self.count_star_inserted = serialized_fn.count_star_inserted
        self.state_cache_size = serialized_fn.state_cache_size
        self.state_cleaning_enabled = serialized_fn.state_cleaning_enabled
        self.data_view_specs = extract_data_view_specs(serialized_fn.udfs)
        super(AbstractStreamGroupAggregateOperation, self).__init__(
            serialized_fn, keyed_state_backend)

    def open(self):
        self.group_agg_function.open(FunctionContext(self.base_metric_group))

    def close(self):
        self.group_agg_function.close()

    def generate_func(self, serialized_fn):
        user_defined_aggs = []
        input_extractors = []
        filter_args = []
        # stores the indexes of the distinct views which the agg functions used
        distinct_indexes = []
        # stores the indexes of the functions which share the same distinct view
        # and the filter args of them
        distinct_info_dict = {}
        for i in range(len(serialized_fn.udfs)):
            user_defined_agg, input_extractor, filter_arg, distinct_index = \
                extract_user_defined_aggregate_function(
                    i, serialized_fn.udfs[i], distinct_info_dict)
            user_defined_aggs.append(user_defined_agg)
            input_extractors.append(input_extractor)
            filter_args.append(filter_arg)
            distinct_indexes.append(distinct_index)
        distinct_view_descriptors = {}
        for agg_index_list, filter_arg_list in distinct_info_dict.values():
            if -1 in filter_arg_list:
                # If there is a non-filter call, we don't need to check filter or not before
                # writing the distinct data view.
                filter_arg_list = []
            # use the agg index of the first function as the key of shared distinct view
            distinct_view_descriptors[agg_index_list[0]] = DistinctViewDescriptor(
                input_extractors[agg_index_list[0]], filter_arg_list)

        key_selector = RowKeySelector(self.grouping)
        if len(self.data_view_specs) > 0:
            state_value_coder = DataViewFilterCoder(self.data_view_specs)
        else:
            state_value_coder = PickleCoder()

        self.group_agg_function = self.create_process_function(
            user_defined_aggs, input_extractors, filter_args, distinct_indexes,
            distinct_view_descriptors, key_selector, state_value_coder)

        return self.process_element_or_timer, []

    def process_element_or_timer(self, input_data: Tuple[int, Row, int, Row]):
        # the structure of the input data:
        # [element_type, element(for process_element), timestamp(for timer), key(for timer)]
        # all the fields are nullable except the "element_type"
        if input_data[0] == NORMAL_RECORD:
            if has_cython:
                row = InternalRow.from_row(input_data[1])
            else:
                row = input_data[1]
            self.group_agg_function.process_element(row)
        else:
            if has_cython:
                timer = InternalRow.from_row(input_data[3])
            else:
                timer = input_data[3]
            self.group_agg_function.on_timer(timer)

    @abc.abstractmethod
    def create_process_function(self, user_defined_aggs, input_extractors, filter_args,
                                distinct_indexes, distinct_view_descriptors, key_selector,
                                state_value_coder):
        pass


class StreamGroupAggregateOperation(AbstractStreamGroupAggregateOperation, BundleOperation):

    def __init__(self, serialized_fn, keyed_state_backend):
        super(StreamGroupAggregateOperation, self).__init__(serialized_fn, keyed_state_backend)

    def finish_bundle(self):
        return self.group_agg_function.finish_bundle()

    def create_process_function(self, user_defined_aggs, input_extractors, filter_args,
                                distinct_indexes, distinct_view_descriptors, key_selector,
                                state_value_coder):
        aggs_handler_function = SimpleAggsHandleFunction(
            user_defined_aggs,
            input_extractors,
            self.index_of_count_star,
            self.count_star_inserted,
            self.data_view_specs,
            filter_args,
            distinct_indexes,
            distinct_view_descriptors)

        return GroupAggFunction(
            aggs_handler_function,
            key_selector,
            self.keyed_state_backend,
            state_value_coder,
            self.generate_update_before,
            self.state_cleaning_enabled,
            self.index_of_count_star)


class StreamGroupTableAggregateOperation(AbstractStreamGroupAggregateOperation, BundleOperation):
    def __init__(self, serialized_fn, keyed_state_backend):
        super(StreamGroupTableAggregateOperation, self).__init__(serialized_fn, keyed_state_backend)

    def finish_bundle(self):
        return self.group_agg_function.finish_bundle()

    def create_process_function(self, user_defined_aggs, input_extractors, filter_args,
                                distinct_indexes, distinct_view_descriptors, key_selector,
                                state_value_coder):
        aggs_handler_function = SimpleTableAggsHandleFunction(
            user_defined_aggs,
            input_extractors,
            self.data_view_specs,
            filter_args,
            distinct_indexes,
            distinct_view_descriptors)
        return GroupTableAggFunction(
            aggs_handler_function,
            key_selector,
            self.keyed_state_backend,
            state_value_coder,
            self.generate_update_before,
            self.state_cleaning_enabled,
            self.index_of_count_star)


class StreamGroupWindowAggregateOperation(AbstractStreamGroupAggregateOperation):
    def __init__(self, serialized_fn, keyed_state_backend):
        self._window = serialized_fn.group_window
        self._named_property_extractor = self._create_named_property_function()
        self._is_time_window = None
        self._reuse_timer_data = Row()
        self._reuse_key_data = Row()
        super(StreamGroupWindowAggregateOperation, self).__init__(
            serialized_fn, keyed_state_backend)

    def create_process_function(self, user_defined_aggs, input_extractors, filter_args,
                                distinct_indexes, distinct_view_descriptors, key_selector,
                                state_value_coder):
        from pyflink.fn_execution import flink_fn_execution_pb2

        self._is_time_window = self._window.is_time_window
        self._namespace_coder = self.keyed_state_backend._namespace_coder_impl
        if self._window.window_type == flink_fn_execution_pb2.GroupWindow.TUMBLING_GROUP_WINDOW:
            if self._is_time_window:
                window_assigner = TumblingWindowAssigner(
                    self._window.window_size, 0, self._window.is_row_time)
            else:
                window_assigner = CountTumblingWindowAssigner(self._window.window_size)
        elif self._window.window_type == flink_fn_execution_pb2.GroupWindow.SLIDING_GROUP_WINDOW:
            if self._is_time_window:
                window_assigner = SlidingWindowAssigner(
                    self._window.window_size, self._window.window_slide, 0,
                    self._window.is_row_time)
            else:
                window_assigner = CountSlidingWindowAssigner(
                    self._window.window_size, self._window.window_slide)
        else:
            window_assigner = SessionWindowAssigner(
                self._window.window_gap, self._window.is_row_time)
        if self._is_time_window:
            if self._window.is_row_time:
                trigger = EventTimeTrigger()
            else:
                trigger = ProcessingTimeTrigger()
        else:
            trigger = CountTrigger(self._window.window_size)

        window_aggregator = SimpleNamespaceAggsHandleFunction(
            user_defined_aggs,
            input_extractors,
            self.index_of_count_star,
            self.count_star_inserted,
            self._named_property_extractor,
            self.data_view_specs,
            filter_args,
            distinct_indexes,
            distinct_view_descriptors)
        return GroupWindowAggFunction(
            self._window.allowedLateness,
            key_selector,
            self.keyed_state_backend,
            state_value_coder,
            window_assigner,
            window_aggregator,
            trigger,
            self._window.time_field_index,
            self._window.shift_timezone)

    def process_element_or_timer(self, input_data: Tuple[int, Row, int, int, Row]):
        if input_data[0] == NORMAL_RECORD:
            self.group_agg_function.process_watermark(input_data[3])
            if has_cython:
                input_row = InternalRow.from_row(input_data[1])
            else:
                input_row = input_data[1]
            result_datas = self.group_agg_function.process_element(input_row)
            for result_data in result_datas:
                yield [NORMAL_RECORD, result_data, None]
            timers = self.group_agg_function.get_timers()
            for timer in timers:
                timer_operand_type = timer[0]  # type: TimerOperandType
                internal_timer = timer[1]  # type: InternalTimer
                window = internal_timer.get_namespace()
                self._reuse_key_data._values = internal_timer.get_key()
                timestamp = internal_timer.get_timestamp()
                encoded_window = self._namespace_coder.encode(window)
                self._reuse_timer_data._values = \
                    [timer_operand_type.value, self._reuse_key_data, timestamp, encoded_window]
                yield [TRIGGER_TIMER, None, self._reuse_timer_data]
        else:
            timestamp = input_data[2]
            timer_data = input_data[4]
            key = list(timer_data[1])
            timer_type = timer_data[0]
            namespace = self._namespace_coder.decode(timer_data[2])
            timer = InternalTimerImpl(timestamp, key, namespace)
            if timer_type == REGISTER_EVENT_TIMER:
                result_datas = self.group_agg_function.on_event_time(timer)
            else:
                result_datas = self.group_agg_function.on_processing_time(timer)
            for result_data in result_datas:
                yield [NORMAL_RECORD, result_data, None]

    def _create_named_property_function(self):
        from pyflink.fn_execution import flink_fn_execution_pb2

        named_property_extractor_array = []
        for named_property in self._window.namedProperties:
            if named_property == flink_fn_execution_pb2.GroupWindow.WINDOW_START:
                named_property_extractor_array.append("value.start")
            elif named_property == flink_fn_execution_pb2.GroupWindow.WINDOW_END:
                named_property_extractor_array.append("value.end")
            elif named_property == flink_fn_execution_pb2.GroupWindow.ROW_TIME_ATTRIBUTE:
                named_property_extractor_array.append("value.end - 1")
            elif named_property == flink_fn_execution_pb2.GroupWindow.PROC_TIME_ATTRIBUTE:
                named_property_extractor_array.append("-1")
            else:
                raise Exception("Unexpected property %s" % named_property)
        named_property_extractor_str = ','.join(named_property_extractor_array)
        if named_property_extractor_str:
            return eval('lambda value: [%s]' % named_property_extractor_str)
        else:
            return None