from pyflink.datastream.stream_execution_environment import StreamExecutionEnvironment, RuntimeExecutionMode
from pyflink.common import Types


class WordCount(object):
    def __init__(self, env):
        self.env = env
        self.env.set_runtime_mode(RuntimeExecutionMode.BATCH)
        self.env.set_parallelism(1)

    def word_count(self, input_data):
        text = self.env.from_collection(input_data)
        word_count = text.flat_map(lambda x: str(x).lower().split()) \
            .map(lambda i: (i, 1), output_type=Types.TUPLE([Types.STRING(), Types.INT()])) \
            .key_by(lambda i: i[0]) \
            .reduce(lambda i, j: (i[0], i[1] + j[1]))
        word_count.print()
        self.env.execute("some-name")

WordCount(StreamExecutionEnvironment.get_execution_environment()).word_count(
    ["ABC is a good boy. ABC works for an XYZ org"])


########################################test case ##############################

from pyflink.datastream import StreamExecutionEnvironment, RuntimeExecutionMode
from services.word_count import WordCount
from servicest.test_utils import PyFlinkStreamingTestCase, DataStreamTestSinkFunction

class WordCountTest(PyFlinkStreamingTestCase):
    def setUp(self) -> None:
        super(WordCountTest, self).setUp()
        self.test_sink = DataStreamTestSinkFunction()

    def tearDown(self) -> None:
        self.test_sink.clear()

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

    def test_word_count(self):
        env = StreamExecutionEnvironment.get_execution_environment()
        env.set_runtime_mode(RuntimeExecutionMode.BATCH)
        env.set_parallelism(1)

        # Input data
        input_data = ["ABC is a good boy. ABC works for an XYZ org"]

        # Expected output data
        expected_output = [('abc', 2), ('an', 1), ('for', 1), ('good', 1), ('is', 1), ('org', 1), ('works', 1),
                           ('boy.', 1), ('xyz', 1)]

        # Create WordCount object
        wc = WordCount(env)

        # Create collect sink
        sink = WordCountTest()

        # Execute word_count function
        wc.word_count(input_data).add_sink(sink)

        # Get results from collect sink
        results = self.test_sink.get_results()

        # Assert that the actual output matches the expected output
        self.assert_equals_sorted(results, expected_output)