################################################################################
#  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 typing import TypeVar, Callable, Optional, List, Generic

from pyflink.datastream.functions import AsyncRetryStrategy, AsyncRetryPredicate

__all__ = [
    'RetryPredicate',
    'NoRetryStrategy',
    'FixedDelayRetryStrategy',
    'ExponentialBackoffDelayRetryStrategy'
]


OUT = TypeVar('OUT')


class RetryPredicate(AsyncRetryPredicate, Generic[OUT]):

    def __init__(self,
                 result_predicate: Optional[Callable[[List[OUT]], bool]],
                 exception_predicate: Optional[Callable[[Exception], bool]]):
        self._result_predicate = result_predicate
        self._exception_predicate = exception_predicate

    def result_predicate(self) -> Optional[Callable[[List[OUT]], bool]]:
        return self._result_predicate

    def exception_predicate(self) -> Optional[Callable[[Exception], bool]]:
        return self._exception_predicate


class NoRetryStrategy(AsyncRetryStrategy, Generic[OUT]):

    def can_retry(self, current_attempts: int) -> bool:
        return False

    def get_backoff_time_millis(self, current_attempts: int) -> int:
        return -1

    def get_retry_predicate(self) -> AsyncRetryPredicate[OUT]:
        return RetryPredicate(None, None)


class FixedDelayRetryStrategy(AsyncRetryStrategy, Generic[OUT]):

    def __init__(self,
                 max_attempts: int,
                 backoff_time_millis: int,
                 result_predicate: Optional[Callable[[List[OUT]], bool]],
                 exception_predicate: Optional[Callable[[Exception], bool]]):
        assert max_attempts > 0, "max_attempts should be greater than zero."
        assert backoff_time_millis > 0, "backoff_time_millis should be greater than zero."
        self._max_attempts = max_attempts
        self._backoff_time_millis = backoff_time_millis
        self._result_predicate = result_predicate
        self._exception_predicate = exception_predicate

    def can_retry(self, current_attempts: int) -> bool:
        return current_attempts <= self._max_attempts

    def get_backoff_time_millis(self, current_attempts: int) -> int:
        return self._backoff_time_millis

    def get_retry_predicate(self) -> AsyncRetryPredicate[OUT]:
        return RetryPredicate(self._result_predicate, self._exception_predicate)


class ExponentialBackoffDelayRetryStrategy(AsyncRetryStrategy, Generic[OUT]):

    def __init__(self,
                 max_attempts: int,
                 initial_delay: int,
                 max_retry_delay: int,
                 multiplier: float,
                 result_predicate: Optional[Callable[[List[OUT]], bool]],
                 exception_predicate: Optional[Callable[[Exception], bool]]):
        assert max_attempts > 0, "max_attempts should be greater than zero."
        assert initial_delay > 0, "initial_delay should be greater than zero."
        assert max_retry_delay > 0, "max_retry_delay should be greater than zero."
        assert multiplier > 0, "multiplier should be greater than zero."
        self._max_attempts = max_attempts
        self._max_retry_delay = max_retry_delay
        self._multiplier = multiplier
        self._result_predicate = result_predicate
        self._exception_predicate = exception_predicate
        self._initial_delay = initial_delay
        self._last_retry_delay = initial_delay

    def can_retry(self, current_attempts: int) -> bool:
        return current_attempts <= self._max_attempts

    def get_backoff_time_millis(self, current_attempts: int) -> int:
        if current_attempts <= 1:
            self._last_retry_delay = self._initial_delay
            return self._last_retry_delay

        backoff = int(min(self._last_retry_delay * self._multiplier, self._max_retry_delay))
        self._last_retry_delay = backoff
        return backoff

    def get_retry_predicate(self) -> AsyncRetryPredicate[OUT]:
        return RetryPredicate(self._result_predicate, self._exception_predicate)


NO_RETRY_STRATEGY: AsyncRetryStrategy = NoRetryStrategy()