################################################################################
#  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 sys
from abc import ABC, abstractmethod
from typing import Generic, TypeVar, List, Iterable

from apache_beam.coders import Coder, PickleCoder

from pyflink.datastream.state import StateDescriptor, State, ValueStateDescriptor, \
    ListStateDescriptor, MapStateDescriptor
from pyflink.datastream.timerservice import InternalTimerService
from pyflink.datastream.window import TimeWindow, CountWindow
from pyflink.fn_execution.timerservice_impl import InternalTimerServiceImpl
from pyflink.fn_execution.internal_state import InternalMergingState
from pyflink.fn_execution.state_impl import RemoteKeyedStateBackend

MAX_LONG_VALUE = sys.maxsize

K = TypeVar('K')
W = TypeVar('W', TimeWindow, CountWindow)


class Context(Generic[K, W], ABC):
    """
    Information available in an invocation of methods of InternalWindowProcessFunction.
    """

    @abstractmethod
    def get_partitioned_state(self, state_descriptor: StateDescriptor) -> State:
        """
        Creates a partitioned state handle, using the state backend configured for this task.
        """
        pass

    @abstractmethod
    def current_key(self) -> K:
        """
        Returns current key of current processed element.
        """
        pass

    @abstractmethod
    def current_processing_time(self) -> int:
        """
        Returns the current processing time.
        """
        pass

    @abstractmethod
    def current_watermark(self) -> int:
        """
        Returns the current event-time watermark.
        """
        pass

    @abstractmethod
    def get_window_accumulators(self, window: W) -> List:
        """
        Gets the accumulators of the given window.
        """
        pass

    @abstractmethod
    def set_window_accumulators(self, window: W, acc: List):
        """
        Sets the accumulators of the given window.
        """
        pass

    @abstractmethod
    def clear_window_state(self, window: W):
        """
        Clear window state of the given window.
        """
        pass

    @abstractmethod
    def clear_trigger(self, window: W):
        """
        Call Trigger#clear(Window) on trigger.
        """
        pass

    @abstractmethod
    def on_merge(self, new_window: W, merged_windows: Iterable[W]):
        """
        Call Trigger.on_merge() on trigger.
        """
        pass

    @abstractmethod
    def delete_cleanup_timer(self, window: W):
        """
        Deletes the cleanup timer set for the contents of the provided window.
        """
        pass


class WindowContext(Context[K, W]):
    """
    Context of window.
    """

    def __init__(self,
                 window_operator,
                 trigger_context: 'TriggerContext',
                 state_backend: RemoteKeyedStateBackend,
                 state_value_coder: Coder,
                 timer_service: InternalTimerServiceImpl,
                 is_event_time: bool):
        self._window_operator = window_operator
        self._trigger_context = trigger_context
        self._state_backend = state_backend
        self.timer_service = timer_service
        self.is_event_time = is_event_time
        self.window_state = self._state_backend.get_value_state("window_state", state_value_coder)

    def get_partitioned_state(self, state_descriptor: StateDescriptor) -> State:
        return self._trigger_context.get_partitioned_state(state_descriptor)

    def current_key(self) -> K:
        return self._state_backend.get_current_key()

    def current_processing_time(self) -> int:
        return self.timer_service.current_processing_time()

    def current_watermark(self) -> int:
        return self.timer_service.current_watermark()

    def get_window_accumulators(self, window: W) -> List:
        self.window_state.set_current_namespace(window)
        return self.window_state.value()

    def set_window_accumulators(self, window: W, acc: List):
        self.window_state.set_current_namespace(window)
        self.window_state.update(acc)

    def clear_window_state(self, window: W):
        self.window_state.set_current_namespace(window)
        self.window_state.clear()

    def clear_trigger(self, window: W):
        self._trigger_context.window = window
        self._trigger_context.clear()

    def on_merge(self, new_window: W, merged_windows: Iterable[W]):
        self._trigger_context.window = new_window
        self._trigger_context.merged_windows = merged_windows
        self._trigger_context.on_merge()

    def delete_cleanup_timer(self, window: W):
        cleanup_time = self._window_operator.cleanup_time(window)
        if cleanup_time == MAX_LONG_VALUE:
            # no need to clean up because we didn't set one
            return
        if self.is_event_time:
            self._trigger_context.delete_event_time_timer(cleanup_time)
        else:
            self._trigger_context.delete_processing_time_timer(cleanup_time)


class TriggerContext(object):
    """
    TriggerContext is a utility for handling Trigger invocations. It can be reused by setting the
    key and window fields. No internal state must be kept in the TriggerContext.
    """

    def __init__(self,
                 trigger,
                 timer_service: InternalTimerService[W],
                 state_backend: RemoteKeyedStateBackend):
        self._trigger = trigger
        self._timer_service = timer_service
        self._state_backend = state_backend
        self.window = None  # type: W
        self.merged_windows = None  # type: Iterable[W]

    def open(self):
        self._trigger.open(self)

    def on_element(self, row, timestamp: int) -> bool:
        return self._trigger.on_element(row, timestamp, self.window)

    def on_processing_time(self, timestamp: int) -> bool:
        return self._trigger.on_processing_time(timestamp, self.window)

    def on_event_time(self, timestamp: int) -> bool:
        return self._trigger.on_event_time(timestamp, self.window)

    def on_merge(self):
        self._trigger.on_merge(self.window, self)

    def get_current_processing_time(self) -> int:
        return self._timer_service.current_processing_time()

    def get_current_watermark(self) -> int:
        return self._timer_service.current_watermark()

    def register_processing_time_timer(self, time: int):
        self._timer_service.register_processing_time_timer(self.window, time)

    def register_event_time_timer(self, time: int):
        self._timer_service.register_event_time_timer(self.window, time)

    def delete_processing_time_timer(self, time: int):
        self._timer_service.delete_processing_time_timer(self.window, time)

    def delete_event_time_timer(self, time: int):
        self._timer_service.delete_event_time_timer(self.window, time)

    def clear(self):
        self._trigger.clear(self.window)

    def get_partitioned_state(self, state_descriptor: StateDescriptor) -> State:
        if isinstance(state_descriptor, ValueStateDescriptor):
            state = self._state_backend.get_value_state(state_descriptor.name, PickleCoder())
        elif isinstance(state_descriptor, ListStateDescriptor):
            state = self._state_backend.get_list_state(state_descriptor.name, PickleCoder())
        elif isinstance(state_descriptor, MapStateDescriptor):
            state = self._state_backend.get_map_state(
                state_descriptor.name, PickleCoder(), PickleCoder())
        else:
            raise Exception("Unknown supported StateDescriptor %s" % state_descriptor)
        state.set_current_namespace(self.window)
        return state

    def merge_partitioned_state(self, state_descriptor: StateDescriptor):
        if not self.merged_windows:
            state = self.get_partitioned_state(state_descriptor)
            if isinstance(state, InternalMergingState):
                state.merge_namespaces(self.window, self.merged_windows)
            else:
                raise Exception("The given state descriptor does not refer to a mergeable state"
                                " (MergingState)")
