################################################################################
#  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.
################################################################################
from pyflink.datastream import OutputTag
from pyflink.datastream.window import WindowOperationDescriptor
from pyflink.fn_execution import pickle
from pyflink.fn_execution.coders import TimeWindowCoder, CountWindowCoder
from pyflink.fn_execution.datastream import operations
from pyflink.fn_execution.datastream.embedded.process_function import (
    InternalProcessFunctionContext, InternalKeyedProcessFunctionContext,
    InternalKeyedProcessFunctionOnTimerContext, InternalWindowTimerContext,
    InternalBroadcastProcessFunctionContext, InternalBroadcastProcessFunctionReadOnlyContext,
    InternalKeyedBroadcastProcessFunctionContext,
    InternalKeyedBroadcastProcessFunctionReadOnlyContext,
    InternalKeyedBroadcastProcessFunctionOnTimerContext)
from pyflink.fn_execution.datastream.embedded.runtime_context import StreamingRuntimeContext
from pyflink.fn_execution.datastream.embedded.side_output_context import SideOutputContext
from pyflink.fn_execution.datastream.embedded.timerservice_impl import InternalTimerServiceImpl
from pyflink.fn_execution.datastream.window.window_operator import WindowOperator
from pyflink.fn_execution.embedded.converters import (TimeWindowConverter, CountWindowConverter,
                                                      GlobalWindowConverter)
from pyflink.fn_execution.embedded.state_impl import KeyedStateBackend


class OneInputOperation(operations.OneInputOperation):

    def __init__(self, open_func, close_func, process_element_func, on_timer_func=None):
        self._open_func = open_func
        self._close_func = close_func
        self._process_element_func = process_element_func
        self._on_timer_func = on_timer_func

    def open(self) -> None:
        self._open_func()

    def close(self) -> None:
        self._close_func()

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

    def on_timer(self, timestamp):
        if self._on_timer_func:
            return self._on_timer_func(timestamp)


class TwoInputOperation(operations.TwoInputOperation):

    def __init__(self, open_func, close_func, process_element_func1, process_element_func2,
                 on_timer_func=None):
        self._open_func = open_func
        self._close_func = close_func
        self._process_element_func1 = process_element_func1
        self._process_element_func2 = process_element_func2
        self._on_timer_func = on_timer_func

    def open(self) -> None:
        self._open_func()

    def close(self) -> None:
        self._close_func()

    def process_element1(self, value):
        return self._process_element_func1(value)

    def process_element2(self, value):
        return self._process_element_func2(value)

    def on_timer(self, timestamp):
        if self._on_timer_func:
            return self._on_timer_func(timestamp)


def extract_process_function(
        user_defined_function_proto, j_runtime_context, j_function_context, j_timer_context,
        j_side_output_context, job_parameters, j_keyed_state_backend, j_operator_state_backend):
    from pyflink.fn_execution import flink_fn_execution_pb2
    UserDefinedDataStreamFunction = flink_fn_execution_pb2.UserDefinedDataStreamFunction

    user_defined_func = pickle.loads(user_defined_function_proto.payload)
    func_type = user_defined_function_proto.function_type

    runtime_context = StreamingRuntimeContext.of(j_runtime_context, job_parameters)

    if j_side_output_context:
        side_output_context = SideOutputContext(j_side_output_context)

        def process_func(values):
            if values is None:
                return
            for value in values:
                if isinstance(value, tuple) and isinstance(value[0], OutputTag):
                    output_tag: OutputTag = value[0]
                    side_output_context.collect(output_tag.tag_id, value[1])
                else:
                    yield value
    else:
        def process_func(values):
            if values is None:
                return
            yield from values

    def open_func():
        if hasattr(user_defined_func, "open"):
            user_defined_func.open(runtime_context)

    def close_func():
        if hasattr(user_defined_func, "close"):
            user_defined_func.close()

    if func_type == UserDefinedDataStreamFunction.PROCESS:
        function_context = InternalProcessFunctionContext(j_function_context)
        process_element = user_defined_func.process_element

        def process_element_func(value):
            yield from process_func(process_element(value, function_context))

        return OneInputOperation(open_func, close_func, process_element_func)

    elif func_type == UserDefinedDataStreamFunction.KEYED_PROCESS:
        function_context = InternalKeyedProcessFunctionContext(
            j_function_context, user_defined_function_proto.key_type_info)
        timer_context = InternalKeyedProcessFunctionOnTimerContext(
            j_timer_context, user_defined_function_proto.key_type_info)

        keyed_state_backend = KeyedStateBackend(
            function_context,
            j_keyed_state_backend)
        runtime_context.set_keyed_state_backend(keyed_state_backend)

        process_element = user_defined_func.process_element
        on_timer = user_defined_func.on_timer

        def process_element_func(value):
            yield from process_func(process_element(value[1], function_context))

        def on_timer_func(timestamp):
            yield from process_func(on_timer(timestamp, timer_context))

        return OneInputOperation(open_func, close_func, process_element_func, on_timer_func)

    elif func_type == UserDefinedDataStreamFunction.CO_PROCESS:
        function_context = InternalProcessFunctionContext(j_function_context)

        process_element1 = user_defined_func.process_element1
        process_element2 = user_defined_func.process_element2

        def process_element_func1(value):
            yield from process_func(process_element1(value, function_context))

        def process_element_func2(value):
            yield from process_func(process_element2(value, function_context))

        return TwoInputOperation(
            open_func, close_func, process_element_func1, process_element_func2)

    elif func_type == UserDefinedDataStreamFunction.CO_BROADCAST_PROCESS:
        broadcast_ctx = InternalBroadcastProcessFunctionContext(
            j_function_context, j_operator_state_backend)
        read_only_broadcast_ctx = InternalBroadcastProcessFunctionReadOnlyContext(
            j_function_context, j_operator_state_backend)

        process_element = user_defined_func.process_element
        process_broadcast_element = user_defined_func.process_broadcast_element

        def process_element_func1(value):
            yield from process_func(process_element(value, read_only_broadcast_ctx))

        def process_element_func2(value):
            yield from process_func(process_broadcast_element(value, broadcast_ctx))

        return TwoInputOperation(
            open_func, close_func, process_element_func1, process_element_func2)

    elif func_type == UserDefinedDataStreamFunction.KEYED_CO_PROCESS:
        function_context = InternalKeyedProcessFunctionContext(
            j_function_context, user_defined_function_proto.key_type_info)
        timer_context = InternalKeyedProcessFunctionOnTimerContext(
            j_timer_context, user_defined_function_proto.key_type_info)

        keyed_state_backend = KeyedStateBackend(
            function_context,
            j_keyed_state_backend)
        runtime_context.set_keyed_state_backend(keyed_state_backend)

        process_element1 = user_defined_func.process_element1
        process_element2 = user_defined_func.process_element2
        on_timer = user_defined_func.on_timer

        def process_element_func1(value):
            yield from process_func(process_element1(value[1], function_context))

        def process_element_func2(value):
            yield from process_func(process_element2(value[1], function_context))

        def on_timer_func(timestamp):
            yield from process_func(on_timer(timestamp, timer_context))

        return TwoInputOperation(
            open_func, close_func, process_element_func1, process_element_func2, on_timer_func)

    elif func_type == UserDefinedDataStreamFunction.KEYED_CO_BROADCAST_PROCESS:
        broadcast_ctx = InternalKeyedBroadcastProcessFunctionContext(
            j_function_context, j_operator_state_backend)
        read_only_broadcast_ctx = InternalKeyedBroadcastProcessFunctionReadOnlyContext(
            j_function_context, user_defined_function_proto.key_type_info, j_operator_state_backend)
        timer_context = InternalKeyedBroadcastProcessFunctionOnTimerContext(
            j_timer_context, user_defined_function_proto.key_type_info, j_operator_state_backend)

        keyed_state_backend = KeyedStateBackend(
            read_only_broadcast_ctx,
            j_keyed_state_backend)
        runtime_context.set_keyed_state_backend(keyed_state_backend)

        process_element = user_defined_func.process_element
        process_broadcast_element = user_defined_func.process_broadcast_element
        on_timer = user_defined_func.on_timer

        def process_element_func1(value):
            yield from process_func(process_element(value[1], read_only_broadcast_ctx))

        def process_element_func2(value):
            yield from process_func(process_broadcast_element(value, broadcast_ctx))

        def on_timer_func(timestamp):
            yield from on_timer(timestamp, timer_context)

        return TwoInputOperation(
            open_func, close_func, process_element_func1, process_element_func2, on_timer_func)

    elif func_type == UserDefinedDataStreamFunction.WINDOW:

        window_operation_descriptor: WindowOperationDescriptor = (
            user_defined_func
        )

        def user_key_selector(normal_data):
            return normal_data

        window_assigner = window_operation_descriptor.assigner
        window_trigger = window_operation_descriptor.trigger
        allowed_lateness = window_operation_descriptor.allowed_lateness
        late_data_output_tag = window_operation_descriptor.late_data_output_tag
        window_state_descriptor = window_operation_descriptor.window_state_descriptor
        internal_window_function = window_operation_descriptor.internal_window_function
        window_serializer = window_operation_descriptor.window_serializer
        window_coder = window_serializer._get_coder()

        if isinstance(window_coder, TimeWindowCoder):
            window_converter = TimeWindowConverter()
        elif isinstance(window_coder, CountWindowCoder):
            window_converter = CountWindowConverter()
        else:
            window_converter = GlobalWindowConverter()

        internal_timer_service = InternalTimerServiceImpl(
            j_timer_context.timerService(), window_converter)

        function_context = InternalKeyedProcessFunctionContext(
            j_function_context,
            user_defined_function_proto.key_type_info)
        window_timer_context = InternalWindowTimerContext(
            j_timer_context,
            user_defined_function_proto.key_type_info,
            window_converter)

        keyed_state_backend = KeyedStateBackend(
            function_context,
            j_keyed_state_backend,
            j_function_context.getWindowSerializer(),
            window_converter)
        runtime_context.set_keyed_state_backend(keyed_state_backend)

        window_operator = WindowOperator(
            window_assigner,
            keyed_state_backend,
            user_key_selector,
            window_state_descriptor,
            internal_window_function,
            window_trigger,
            allowed_lateness,
            late_data_output_tag)

        def open_func():
            window_operator.open(runtime_context, internal_timer_service)

        def close_func():
            window_operator.close()

        def process_element_func(value):
            yield from process_func(
                window_operator.process_element(value[1], function_context.timestamp()))

        if window_assigner.is_event_time():
            def on_timer_func(timestamp):
                window = window_timer_context.window()
                key = window_timer_context.get_current_key()
                yield from process_func(window_operator.on_event_time(timestamp, key, window))
        else:
            def on_timer_func(timestamp):
                window = window_timer_context.window()
                key = window_timer_context.get_current_key()
                yield from process_func(window_operator.on_processing_time(timestamp, key, window))

        return OneInputOperation(open_func, close_func, process_element_func, on_timer_func)

    else:
        raise Exception("Unknown function type {0}.".format(func_type))
