################################################################################
#  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 asyncio
import random
from typing import List

from pyflink.common import Types, Row, Time, Configuration, WatermarkStrategy
from pyflink.datastream import AsyncDataStream, AsyncFunction, \
    StreamExecutionEnvironment, AsyncRetryStrategy
from pyflink.datastream.tests.test_util import DataStreamTestSinkFunction, \
    SecondColumnTimestampAssigner
from pyflink.java_gateway import get_gateway
from pyflink.testing.test_case_utils import PyFlinkStreamingTestCase
from pyflink.util.java_utils import get_j_env_configuration


class AsyncFunctionTests(PyFlinkStreamingTestCase):

    def setUp(self) -> None:
        super(AsyncFunctionTests, self).setUp()
        config = get_j_env_configuration(self.env._j_stream_execution_environment)
        config.setString("pekko.ask.timeout", "20 s")
        self.test_sink = DataStreamTestSinkFunction()

    def assert_equals_sorted(self, expected, actual):
        expected.sort()
        actual.sort()
        self.assertEqual(expected, actual)

    def assert_equals(self, expected, actual):
        self.assertEqual(expected, actual)

    def test_unordered_mode(self):
        self.env.set_parallelism(1)
        ds = self.env.from_collection(
            [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)],
            type_info=Types.ROW_NAMED(["v1", "v2"], [Types.INT(), Types.INT()])
        )

        class MyAsyncFunction(AsyncFunction):

            async def async_invoke(self, value: Row):
                await asyncio.sleep(2)
                return [value[0] + value[1]]

            def timeout(self, value: Row):
                return [value[0] + value[1]]

        ds = AsyncDataStream.unordered_wait(
            ds, MyAsyncFunction(), Time.seconds(5), 2, Types.INT())
        ds.add_sink(self.test_sink)
        self.env.execute()
        results = self.test_sink.get_results(False)
        expected = ['2', '4', '6', '8', '10']
        self.assert_equals_sorted(expected, results)

    def test_ordered_mode(self):
        self.env.set_parallelism(1)
        ds = self.env.from_collection(
            [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)],
            type_info=Types.ROW_NAMED(["v1", "v2"], [Types.INT(), Types.INT()])
        )

        class MyAsyncFunction(AsyncFunction):

            async def async_invoke(self, value: Row):
                await asyncio.sleep(random.randint(1, 2))
                return [value[0] + value[1]]

            def timeout(self, value: Row):
                return [value[0] + value[1]]

        ds = AsyncDataStream.ordered_wait(
            ds, MyAsyncFunction(), Time.seconds(5), 2, Types.INT())
        ds.add_sink(self.test_sink)
        self.env.execute()
        results = self.test_sink.get_results(False)
        expected = ['2', '4', '6', '8', '10']
        self.assert_equals(expected, results)

    def test_watermark(self):
        self.env.set_parallelism(1)
        ds = self.env.from_collection(
            [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)],
            type_info=Types.ROW_NAMED(["v1", "v2"], [Types.INT(), Types.INT()])
        )
        jvm = get_gateway().jvm
        watermark_strategy = WatermarkStrategy(
            jvm.org.apache.flink.api.common.eventtime.WatermarkStrategy.forGenerator(
                jvm.org.apache.flink.streaming.api.functions.python.eventtime.
                PerElementWatermarkGenerator.getSupplier()
            )
        ).with_timestamp_assigner(SecondColumnTimestampAssigner())
        ds = ds.assign_timestamps_and_watermarks(watermark_strategy)

        class MyAsyncFunction(AsyncFunction):

            async def async_invoke(self, value: Row):
                await asyncio.sleep(random.randint(1, 3))
                return [value[0] + value[1]]

            def timeout(self, value: Row):
                return [value[0] + value[1]]

        ds = AsyncDataStream.unordered_wait(
            ds, MyAsyncFunction(), Time.seconds(5), 2, Types.INT())
        ds.add_sink(self.test_sink)
        self.env.execute()
        results = self.test_sink.get_results(False)
        expected = ['2', '4', '6', '8', '10']
        # note that we use assert_equals instead of assert_equals_sorted
        self.assert_equals(expected, results)

    def test_non_iterable_result(self):
        self.env.set_parallelism(1)
        ds = self.env.from_collection(
            [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)],
            type_info=Types.ROW_NAMED(["v1", "v2"], [Types.INT(), Types.INT()])
        )

        class MyAsyncFunction(AsyncFunction):

            async def async_invoke(self, value: Row):
                await asyncio.sleep(2)
                return value[0] + value[1]

            def timeout(self, value: Row):
                return value[0] + value[1]

        ds = AsyncDataStream.unordered_wait(
            ds, MyAsyncFunction(), Time.seconds(5), 2, Types.INT())
        ds.add_sink(self.test_sink)
        try:
            self.env.execute()
            self.fail()
        except Exception as e:
            message = str(e)
            self.assertTrue("The result of AsyncFunction should be of list type" in message)

    def test_none_result(self):
        self.env.set_parallelism(1)
        ds = self.env.from_collection(
            [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)],
            type_info=Types.ROW_NAMED(["v1", "v2"], [Types.INT(), Types.INT()])
        )

        class MyAsyncFunction(AsyncFunction):

            async def async_invoke(self, value: Row):
                await asyncio.sleep(10)
                return None

            def timeout(self, value: Row):
                return None

        ds = AsyncDataStream.unordered_wait(
            ds, MyAsyncFunction(), Time.seconds(1), 2, Types.INT())
        ds.add_sink(self.test_sink)
        try:
            self.env.execute()
            self.fail()
        except Exception as e:
            message = str(e)
            self.assertTrue("The result of AsyncFunction cannot be none" in message)

    def test_raise_exception_in_async_invoke(self):
        self.env.set_parallelism(1)
        ds = self.env.from_collection(
            [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)],
            type_info=Types.ROW_NAMED(["v1", "v2"], [Types.INT(), Types.INT()])
        )

        class MyAsyncFunction(AsyncFunction):

            async def async_invoke(self, value: Row):
                raise Exception("encountered an exception")

            def timeout(self, value: Row):
                # raise the same exception to make sure test case is stable in all cases
                raise Exception("encountered an exception")

        ds = AsyncDataStream.unordered_wait(
            ds, MyAsyncFunction(), Time.seconds(5), 2, Types.INT())
        ds.add_sink(self.test_sink)
        try:
            self.env.execute()
            self.fail()
        except Exception as e:
            message = str(e)
            self.assertTrue("encountered an exception" in message)

    def test_raise_exception_in_timeout(self):
        self.env.set_parallelism(1)
        ds = self.env.from_collection(
            [(1, 1), (2, 2), (3, 3)],
            type_info=Types.ROW_NAMED(["v1", "v2"], [Types.INT(), Types.INT()])
        )

        class MyAsyncFunction(AsyncFunction):

            async def async_invoke(self, value: Row):
                await asyncio.sleep(10)
                return [value[0] + value[1]]

            def timeout(self, value: Row):
                raise Exception("encountered an exception")

        ds = AsyncDataStream.unordered_wait(
            ds, MyAsyncFunction(), Time.seconds(2), 2, Types.INT())
        ds.add_sink(self.test_sink)
        try:
            self.env.execute()
            self.fail()
        except Exception as e:
            message = str(e)
            self.assertTrue("encountered an exception" in message)

    def test_processing_timeout(self):
        self.env.set_parallelism(1)
        ds = self.env.from_collection(
            [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)],
            type_info=Types.ROW_NAMED(["v1", "v2"], [Types.INT(), Types.INT()])
        )

        class MyAsyncFunction(AsyncFunction):

            async def async_invoke(self, value: Row):
                await asyncio.sleep(10)
                return [value[0] + value[1]]

            def timeout(self, value: Row):
                return [value[0] - value[1]]

        ds = AsyncDataStream.unordered_wait(
            ds, MyAsyncFunction(), Time.seconds(1), 2, Types.INT())
        ds.add_sink(self.test_sink)
        self.env.execute()
        results = self.test_sink.get_results(False)
        expected = ['0', '0', '0', '0', '0']
        self.assert_equals_sorted(expected, results)

    def test_async_with_retry(self):
        self.env.set_parallelism(1)
        ds = self.env.from_collection(
            [(1, 1), (2, 2), (3, 3)],
            type_info=Types.ROW_NAMED(["v1", "v2"], [Types.INT(), Types.INT()])
        )

        class MyAsyncFunction(AsyncFunction):

            def __init__(self):
                self.retries_1 = {}
                self.retries_2 = {}

            async def async_invoke(self, value: Row):
                await asyncio.sleep(1)
                if value in self.retries_2:
                    return [value[0] + value[1]]
                elif value in self.retries_1:
                    self.retries_2[value] = True
                    return [value[0] + value[1] + 1]
                else:
                    self.retries_1[value] = True
                    raise ValueError("failed the first time")

            def timeout(self, value: Row):
                return [value[0] + value[1]]

        def result_predicate(result: List[int]):
            return result[0] % 2 == 1

        def exception_predicate(exception: Exception):
            return "failed the first time" in str(exception)

        async_retry_strategy = AsyncRetryStrategy.fixed_delay(
            max_attempts=5,
            backoff_time_millis=1000,
            result_predicate=result_predicate,
            exception_predicate=exception_predicate
        )
        ds = AsyncDataStream.unordered_wait_with_retry(
            ds, MyAsyncFunction(), Time.seconds(10), async_retry_strategy, 2, Types.INT())
        ds.add_sink(self.test_sink)
        self.env.execute()
        results = self.test_sink.get_results(False)
        expected = ['2', '4', '6']
        self.assert_equals_sorted(expected, results)


class EmbeddedThreadAsyncFunctionTests(PyFlinkStreamingTestCase):

    def test_run_async_function_in_thread_mode(self):
        config = Configuration()
        config.set_string("python.execution-mode", "thread")
        env = StreamExecutionEnvironment.get_execution_environment(config)
        ds = env.from_collection(
            [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)],
            type_info=Types.ROW_NAMED(["v1", "v2"], [Types.INT(), Types.INT()])
        )

        class MyAsyncFunction(AsyncFunction):

            async def async_invoke(self, value: Row):
                await asyncio.sleep(2)
                return [value[0] + value[1]]

        try:
            AsyncDataStream.unordered_wait(
                ds, MyAsyncFunction(), Time.seconds(5), 2, Types.INT())
            self.fail()
        except Exception as e:
            message = str(e)
            self.assertTrue("AsyncFunction is still not supported for 'thread' mode" in message)