#
# 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 os
import random
import tempfile
import unittest

import pandas as pd
import pandas.testing as pdt

from pyspark import SparkConf
from pyspark.sql import DataFrame
from pyspark.sql.functions import split
from pyspark.sql.streaming import StatefulProcessor, TwsTester
from pyspark.sql.streaming.query import StreamingQuery
from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
    AllMethodsTestProcessorFactory,
    EventTimeCountProcessorFactory,
    EventTimeSessionProcessorFactory,
    RunningCountStatefulProcessorFactory,
    SessionTimeoutProcessorFactory,
    TopKProcessorFactory,
    WordFrequencyProcessorFactory,
)
from pyspark.sql.types import (
    DoubleType,
    IntegerType,
    Row,
    StringType,
    StructField,
    StructType,
)
from pyspark.errors import PySparkValueError, PySparkAssertionError
from pyspark.errors.exceptions.base import IllegalArgumentException
from pyspark.testing.sqlutils import ReusedSQLTestCase
from pyspark.testing.utils import (
    have_pandas,
    have_pyarrow,
    pandas_requirement_message,
    pyarrow_requirement_message,
)


@unittest.skipIf(
    not have_pandas or not have_pyarrow or os.environ.get("PYTHON_GIL", "?") == "0",
    pandas_requirement_message or pyarrow_requirement_message or "Not supported in no-GIL mode",
)
class TwsTesterTests(ReusedSQLTestCase):
    def test_running_count_processor(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)

        self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
        self.assertEqual(
            tester.test("key2", [Row(value="a"), Row(value="a")]),
            [Row(key="key2", count=2)],
        )
        self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
        self.assertEqual(
            tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
            [Row(key="key1", count=4)],
        )

        self.assertEqual(tester.peekValueState("count", "key1"), (4,))
        self.assertEqual(tester.peekValueState("count", "key2"), (2,))
        self.assertEqual(tester.peekValueState("count", "key3"), (1,))
        self.assertIsNone(tester.peekValueState("count", "key4"))

    def test_running_count_processor_pandas(self):
        processor = RunningCountStatefulProcessorFactory().pandas()
        tester = TwsTester(processor)

        ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
        expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
        pdt.assert_frame_equal(ans1, expected1, check_like=True)

        ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
        expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
        pdt.assert_frame_equal(ans2, expected2, check_like=True)

        ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
        expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
        pdt.assert_frame_equal(ans3, expected3, check_like=True)

        ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
        expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
        pdt.assert_frame_equal(ans4, expected4, check_like=True)

    def test_direct_access_to_value_state(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)
        tester.updateValueState("count", "foo", (5,))
        tester.test("foo", [Row(value="q")])
        self.assertEqual(tester.peekValueState("count", "foo"), (6,))

    def test_topk_processor(self):
        processor = TopKProcessorFactory(k=2).row()
        tester = TwsTester(processor)

        ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
        self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])

        ans2 = tester.test(
            "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
        )
        self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])

        ans3 = tester.test("key3", [Row(score=100.0)])
        self.assertEqual(ans3, [Row(key="key3", score=100.0)])

        self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
        self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
        self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
        self.assertEqual(tester.peekListState("topK", "key4"), [])

        ans4 = tester.test("key1", [Row(score=10.0)])
        self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
        self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])

    def test_topk_processor_pandas(self):
        processor = TopKProcessorFactory(k=2).pandas()
        tester = TwsTester(processor)

        ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
        expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
        pdt.assert_frame_equal(ans1, expected1, check_like=True)

        ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
        expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
        pdt.assert_frame_equal(ans2, expected2, check_like=True)

        ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
        expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
        pdt.assert_frame_equal(ans3, expected3, check_like=True)

        ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
        expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
        pdt.assert_frame_equal(ans4, expected4, check_like=True)

    def test_direct_access_to_list_state(self):
        processor = TopKProcessorFactory(k=2).row()
        tester = TwsTester(processor)

        tester.updateListState("topK", "a", [(6.0,), (5.0,)])
        tester.updateListState("topK", "b", [(8.0,), (7.0,)])
        tester.test("a", [Row(score=10.0)])
        tester.test("b", [Row(score=7.5)])
        tester.test("c", [Row(score=1.0)])

        assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
        assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
        assert tester.peekListState("topK", "c") == [(1.0,)]
        assert tester.peekListState("topK", "d") == []

    def test_word_frequency_processor(self):
        processor = WordFrequencyProcessorFactory().row()
        tester = TwsTester(processor)

        ans1 = tester.test(
            "user1",
            [
                Row(word="hello"),
                Row(word="world"),
                Row(word="hello"),
                Row(word="world"),
            ],
        )
        self.assertEqual(
            ans1,
            [
                Row(key="user1", word="hello", count=1),
                Row(key="user1", word="world", count=1),
                Row(key="user1", word="hello", count=2),
                Row(key="user1", word="world", count=2),
            ],
        )

        ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
        self.assertEqual(
            ans2,
            [
                Row(key="user2", word="hello", count=1),
                Row(key="user2", word="spark", count=1),
            ],
        )

        # Check state using peekMapState.
        self.assertEqual(
            tester.peekMapState("frequencies", "user1"),
            {("hello",): (2,), ("world",): (2,)},
        )
        self.assertEqual(
            tester.peekMapState("frequencies", "user2"),
            {("hello",): (1,), ("spark",): (1,)},
        )
        self.assertEqual(tester.peekMapState("frequencies", "user3"), {})

        # Process more data for user1.
        ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
        self.assertEqual(
            ans3,
            [
                Row(key="user1", word="hello", count=3),
                Row(key="user1", word="test", count=1),
            ],
        )
        self.assertEqual(
            tester.peekMapState("frequencies", "user1"),
            {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
        )

    def test_word_frequency_processor_pandas(self):
        processor = WordFrequencyProcessorFactory().pandas()
        tester = TwsTester(processor)

        input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
        ans1 = tester.testInPandas("user1", input_df1)
        expected1 = pd.DataFrame(
            {
                "key": ["user1", "user1", "user1", "user1"],
                "word": ["hello", "world", "hello", "world"],
                "count": [1, 1, 2, 2],
            }
        )
        pdt.assert_frame_equal(ans1, expected1, check_like=True)

        input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
        ans2 = tester.testInPandas("user2", input_df2)
        expected2 = pd.DataFrame(
            {
                "key": ["user2", "user2"],
                "word": ["hello", "spark"],
                "count": [1, 1],
            }
        )
        pdt.assert_frame_equal(ans2, expected2, check_like=True)

        input_df3 = pd.DataFrame({"word": ["hello", "test"]})
        ans3 = tester.testInPandas("user1", input_df3)
        expected3 = pd.DataFrame(
            {
                "key": ["user1", "user1"],
                "word": ["hello", "test"],
                "count": [3, 1],
            }
        )
        pdt.assert_frame_equal(ans3, expected3, check_like=True)

    def test_direct_access_to_map_state(self):
        processor = WordFrequencyProcessorFactory().row()
        tester = TwsTester(processor)

        tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
        tester.updateMapState("frequencies", "user2", {("spark",): (10,)})

        tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
        tester.test("user2", [Row(word="spark")])
        tester.test("user3", [Row(word="new")])

        self.assertEqual(
            tester.peekMapState("frequencies", "user1"),
            {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
        )
        self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
        self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
        self.assertEqual(tester.peekMapState("frequencies", "user4"), {})

    # Example of how TwsTester can be used to test step function.
    def test_step_function(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)

        # Example of helper function using TwsTester to inspect how processing a single row changes
        # state.
        def step_function(key: str, input_row: str, state_in: int) -> int:
            tester.updateValueState("count", key, (state_in,))
            tester.test(key, [Row(value=input_row)])
            return tester.peekValueState("count", key)[0]

        self.assertEqual(step_function("key1", "a", 10), 11)

    # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
    def test_row_by_row(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)

        # Example of helper function to test how TransformWithState processes rows one-by-one,
        # which can be used to simulate real-time mode.
        def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
            result: list[Row] = []
            for key, row in input_rows:
                result += tester.test(key, [row])
            return result

        ans = test_row_by_row_helper(
            [
                ("key1", Row(value="a")),
                ("key2", Row(value="b")),
                ("key1", Row(value="c")),
                ("key2", Row(value="b")),
                ("key1", Row(value="c")),
                ("key1", Row(value="c")),
                ("key3", Row(value="q")),
            ]
        )
        self.assertEqual(
            ans,
            [
                Row(key="key1", count=1),
                Row(key="key2", count=1),
                Row(key="key1", count=2),
                Row(key="key2", count=2),
                Row(key="key1", count=3),
                Row(key="key1", count=4),
                Row(key="key3", count=1),
            ],
        )

    # Tests that TwsTester calls handleInitialState.
    def test_initial_state_row(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(
            processor,
            initialStateRow=[
                ("a", Row(initial_count=10)),
                ("b", Row(initial_count=20)),
            ],
        )
        self.assertEqual(tester.peekValueState("count", "a"), (10,))
        self.assertEqual(tester.peekValueState("count", "b"), (20,))

        ans1 = tester.test("a", [Row(value="a")])
        self.assertEqual(ans1, [Row(key="a", count=11)])

        ans2 = tester.test("c", [Row(value="c")])
        self.assertEqual(ans2, [Row(key="c", count=1)])

    def test_initial_state_pandas(self):
        processor = RunningCountStatefulProcessorFactory().pandas()
        tester = TwsTester(
            processor,
            initialStatePandas=[
                ("a", pd.DataFrame({"initial_count": [10]})),
                ("b", pd.DataFrame({"initial_count": [20]})),
            ],
        )
        self.assertEqual(tester.peekValueState("count", "a"), (10,))
        self.assertEqual(tester.peekValueState("count", "b"), (20,))

        ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
        expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
        pdt.assert_frame_equal(ans1, expected1, check_like=True)

        ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
        expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
        pdt.assert_frame_equal(ans2, expected2, check_like=True)

    def test_all_methods_processor(self):
        """Test that TwsTester exercises all state methods."""
        processor = AllMethodsTestProcessorFactory().row()
        tester = TwsTester(processor)

        results = tester.test(
            "k",
            [
                Row(cmd="value-exists"),  # false
                Row(cmd="value-set"),  # set to 42
                Row(cmd="value-exists"),  # true
                Row(cmd="value-clear"),  # clear
                Row(cmd="value-exists"),  # false again
                Row(cmd="list-exists"),  # false
                Row(cmd="list-append"),  # append a, b
                Row(cmd="list-exists"),  # true
                Row(cmd="list-append-array"),  # append c, d
                Row(cmd="list-get"),  # a,b,c,d
                Row(cmd="map-exists"),  # false
                Row(cmd="map-add"),  # add x=1, y=2, z=3
                Row(cmd="map-exists"),  # true
                Row(cmd="map-keys"),  # x,y,z
                Row(cmd="map-values"),  # 1,2,3
                Row(cmd="map-iterator"),  # x=1,y=2,z=3
                Row(cmd="map-remove"),  # remove y
                Row(cmd="map-keys"),  # x,z
                Row(cmd="map-clear"),  # clear map
                Row(cmd="map-exists"),  # false
            ],
        )

        self.assertEqual(
            results,
            [
                Row(key="k", result="value-exists:False"),
                Row(key="k", result="value-set:done"),
                Row(key="k", result="value-exists:True"),
                Row(key="k", result="value-clear:done"),
                Row(key="k", result="value-exists:False"),
                Row(key="k", result="list-exists:False"),
                Row(key="k", result="list-append:done"),
                Row(key="k", result="list-exists:True"),
                Row(key="k", result="list-append-array:done"),
                Row(key="k", result="list-get:a,b,c,d"),
                Row(key="k", result="map-exists:False"),
                Row(key="k", result="map-add:done"),
                Row(key="k", result="map-exists:True"),
                Row(key="k", result="map-keys:x,y,z"),
                Row(key="k", result="map-values:1,2,3"),
                Row(key="k", result="map-iterator:x=1,y=2,z=3"),
                Row(key="k", result="map-remove:done"),
                Row(key="k", result="map-keys:x,z"),
                Row(key="k", result="map-clear:done"),
                Row(key="k", result="map-exists:False"),
            ],
        )

    def test_both_initial_states_specified(self):
        processor = RunningCountStatefulProcessorFactory().row()
        with self.assertRaises(AssertionError):
            TwsTester(
                processor,
                initialStateRow=[("a", Row(initial_count=10))],
                initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
            )

    def test_timer_registration_raises_error_in_none_mode(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor, timeMode="None")
        with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
            tester.handle.registerTimer(12345)

    def test_delete_timer_raises_error_in_none_mode(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor, timeMode="None")
        with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
            tester.handle.deleteTimer(12345)

    def test_list_timers_raises_error_in_none_mode(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor, timeMode="None")
        with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
            list(tester.handle.listTimers())

    def test_empty_input_row(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)
        result = tester.test("key1", [])
        self.assertEqual(result, [Row(key="key1", count=0)])

    def test_empty_input_pandas(self):
        processor = RunningCountStatefulProcessorFactory().pandas()
        tester = TwsTester(processor)
        input_df = pd.DataFrame({"value": []})
        result = tester.testInPandas("key1", input_df)
        expected = pd.DataFrame({"key": ["key1"], "count": [0]})
        pdt.assert_frame_equal(result, expected, check_like=True)

    def test_empty_initial_state_row(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor, initialStateRow=[])
        result = tester.test("key1", [Row(value="a")])
        self.assertEqual(result, [Row(key="key1", count=1)])

    def test_empty_initial_state_pandas(self):
        processor = RunningCountStatefulProcessorFactory().pandas()
        tester = TwsTester(processor, initialStatePandas=[])
        result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
        expected = pd.DataFrame({"key": ["key1"], "count": [1]})
        pdt.assert_frame_equal(result, expected, check_like=True)

    def test_none_values_in_input(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)
        result = tester.test("key1", [Row(value=None), Row(value="a")])
        self.assertEqual(result, [Row(key="key1", count=2)])

    def test_single_row_input(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)
        result = tester.test("key1", [Row(value="a")])
        self.assertEqual(result, [Row(key="key1", count=1)])

    def test_peek_nonexistent_state(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)
        result = tester.peekValueState("count", "nonexistent_key")
        self.assertIsNone(result)

    def test_update_nonexistent_state(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)
        with self.assertRaises(AssertionError):
            tester.updateValueState("nonexistent_state", "key1", (5,))

    def test_peek_state_before_initialization(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)
        with self.assertRaises(AssertionError):
            tester.peekValueState("nonexistent_state", "key1")

    def test_clear_nonexistent_value_state(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)
        tester.handle.setGroupingKey("key1")
        value_state = tester.handle.getValueState("count", "int")
        value_state.clear()

    def test_clear_nonexistent_list_state(self):
        processor = TopKProcessorFactory(k=2).row()
        tester = TwsTester(processor)
        tester.handle.setGroupingKey("key1")
        list_state = tester.handle.getListState("topK", "double")
        list_state.clear()

    def test_clear_nonexistent_map_state(self):
        processor = WordFrequencyProcessorFactory().row()
        tester = TwsTester(processor)
        tester.handle.setGroupingKey("key1")
        map_state = tester.handle.getMapState("frequencies", "string", "int")
        map_state.clear()

    def test_delete_if_exists_nonexistent(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)
        tester.handle.deleteIfExists("nonexistent_state")

    def test_numeric_key(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)
        self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
        self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])

    def test_tuple_key(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)
        result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
        self.assertEqual(result, [Row(key=(1, 2), count=2)])

    def test_special_characters_in_key(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)
        result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
        self.assertEqual(result, [Row(key="key@#$%", count=2)])

    def test_null_key_value_row(self):
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)
        result = tester.test(None, [Row(value="a"), Row(value="b")])
        self.assertEqual(result, [Row(key=None, count=2)])

    def test_null_key_value_pandas(self):
        processor = RunningCountStatefulProcessorFactory().pandas()
        tester = TwsTester(processor)
        input_df = pd.DataFrame({"value": ["a", "b"]})
        result = tester.testInPandas(None, input_df)
        expected = pd.DataFrame({"key": [None], "count": [2]})
        pdt.assert_frame_equal(result, expected, check_like=True)

    def test_pandas_with_multiple_rows(self):
        processor = RunningCountStatefulProcessorFactory().pandas()
        tester = TwsTester(processor)
        input_df = pd.DataFrame({"value": ["a", "b", "c"]})
        result = tester.testInPandas("key1", input_df)
        expected = pd.DataFrame({"key": ["key1"], "count": [3]})
        pdt.assert_frame_equal(result, expected, check_like=True)

    def test_pandas_dtype_preservation(self):
        processor = RunningCountStatefulProcessorFactory().pandas()
        tester = TwsTester(processor)
        input_df = pd.DataFrame({"value": ["a"]})
        result = tester.testInPandas("key1", input_df)
        self.assertEqual(result["count"].dtype, "int64")

    def test_processor_init_called_once(self):
        from pyspark.sql.streaming.stateful_processor import (
            StatefulProcessor,
            StatefulProcessorHandle,
        )
        from typing import Iterator

        init_call_count = [0]

        class InitCountingProcessor(StatefulProcessor):
            def init(self, handle: StatefulProcessorHandle) -> None:
                init_call_count[0] += 1
                self.handle = handle
                self.state = handle.getValueState("count", "int")

            def handleInputRows(self, key, rows, timerValues) -> Iterator:
                val = self.state.get() if self.state.exists() else None
                count = val[0] if val else 0
                for row in rows:
                    count += 1
                self.state.update((count,))
                yield Row(key=key[0], count=count)

        processor = InitCountingProcessor()
        tester = TwsTester(processor)
        tester.test("key1", [Row(value="a")])
        tester.test("key1", [Row(value="b")])
        self.assertEqual(init_call_count[0], 1)

    def test_delete_value_state(self):
        """Test that deleteState correctly deletes value state for a given key."""
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)

        tester.updateValueState("count", "key1", (10,))
        tester.updateValueState("count", "key2", (20,))
        self.assertEqual(tester.peekValueState("count", "key1"), (10,))

        tester.deleteState("count", "key1")
        self.assertIsNone(tester.peekValueState("count", "key1"))
        self.assertEqual(tester.peekValueState("count", "key2"), (20,))

    def test_delete_list_state(self):
        """Test that deleteState correctly deletes list state for a given key."""
        processor = TopKProcessorFactory(k=3).row()
        tester = TwsTester(processor)

        tester.updateListState("topK", "key1", [(1.0,), (2.0,), (3.0,)])
        tester.updateListState("topK", "key2", [(4.0,), (5.0,)])
        self.assertEqual(tester.peekListState("topK", "key1"), [(1.0,), (2.0,), (3.0,)])

        tester.deleteState("topK", "key1")
        self.assertEqual(tester.peekListState("topK", "key1"), [])
        self.assertEqual(tester.peekListState("topK", "key2"), [(4.0,), (5.0,)])

    def test_delete_map_state(self):
        """Test that deleteState correctly deletes map state for a given key."""
        processor = WordFrequencyProcessorFactory().row()
        tester = TwsTester(processor)

        tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
        tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
        self.assertEqual(
            tester.peekMapState("frequencies", "user1"),
            {("hello",): (5,), ("world",): (3,)},
        )

        tester.deleteState("frequencies", "user1")
        self.assertEqual(tester.peekMapState("frequencies", "user1"), {})
        self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (10,)})

    def test_delete_nonexistent_state_raises_error(self):
        """Test that deleteState raises an error for non-existent state."""
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor)

        with self.assertRaisesRegex(PySparkAssertionError, "STATE_NOT_EXISTS"):
            tester.deleteState("nonexistent_state", "key1")

    # Timer tests.

    def test_processing_time_timers(self):
        """Test that TwsTester supports ProcessingTime timers."""
        processor = SessionTimeoutProcessorFactory().row()
        tester = TwsTester(processor, timeMode="ProcessingTime")

        # Process input for key1 - should register a timer at t=10000.
        result1 = tester.test("key1", [Row(value="hello")])
        self.assertEqual(result1, [Row(key="key1", result="received:hello")])

        # Set processing time to 5000 - timer should NOT fire yet.
        expired1 = tester.setProcessingTime(5000)
        self.assertEqual(expired1, [])

        # Process input for key2 at t=5000 - should register timer at t=15000.
        result2 = tester.test("key2", [Row(value="world")])
        self.assertEqual(result2, [Row(key="key2", result="received:world")])

        # Set processing time to 11000 - key1's timer should fire.
        expired2 = tester.setProcessingTime(11000)
        self.assertEqual(expired2, [Row(key="key1", result="session-expired")])

        # Set processing time to 16000 - key2's timer should fire.
        expired3 = tester.setProcessingTime(16000)
        self.assertEqual(expired3, [Row(key="key2", result="session-expired")])

        # Verify state is cleared after session expiry.
        self.assertIsNone(tester.peekValueState("lastSeen", "key1"))
        self.assertIsNone(tester.peekValueState("lastSeen", "key2"))

    def test_processing_time_timers_pandas(self):
        """Test that TwsTester supports ProcessingTime timers in Pandas mode."""
        processor = SessionTimeoutProcessorFactory().pandas()
        tester = TwsTester(processor, timeMode="ProcessingTime")

        # Process input for key1 - should register a timer at t=10000.
        result1 = tester.testInPandas("key1", pd.DataFrame({"value": ["hello"]}))
        expected1 = pd.DataFrame({"key": ["key1"], "result": ["received:hello"]})
        pdt.assert_frame_equal(result1, expected1, check_like=True)

        # Set processing time to 5000 - timer should NOT fire yet.
        expired1 = tester.setProcessingTime(5000)
        self.assertEqual(len(expired1), 0)

        # Process input for key2 at t=5000 - should register timer at t=15000.
        result2 = tester.testInPandas("key2", pd.DataFrame({"value": ["world"]}))
        expected2 = pd.DataFrame({"key": ["key2"], "result": ["received:world"]})
        pdt.assert_frame_equal(result2, expected2, check_like=True)

        # Set processing time to 11000 - key1's timer should fire.
        expired2 = tester.setProcessingTime(11000)
        expected_expired2 = pd.DataFrame({"key": ["key1"], "result": ["session-expired"]})
        pdt.assert_frame_equal(expired2, expected_expired2, check_like=True)

        # Set processing time to 16000 - key2's timer should fire.
        expired3 = tester.setProcessingTime(16000)
        expected_expired3 = pd.DataFrame({"key": ["key2"], "result": ["session-expired"]})
        pdt.assert_frame_equal(expired3, expected_expired3, check_like=True)

        # Verify state is cleared after session expiry.
        self.assertIsNone(tester.peekValueState("lastSeen", "key1"))
        self.assertIsNone(tester.peekValueState("lastSeen", "key2"))

    def test_event_time_timers_manual_watermark(self):
        """Test that TwsTester supports EventTime timers fired by manual watermark advance."""
        processor = EventTimeSessionProcessorFactory().row()

        def event_time_extractor(row: Row) -> int:
            return row.event_time_ms

        tester = TwsTester(
            processor,
            timeMode="EventTime",
            eventTimeExtractor=event_time_extractor,
        )

        # Process event at t=10000 for key1 - registers timer at t=15000.
        result1 = tester.test("key1", [Row(event_time_ms=10000, value="hello")])
        self.assertEqual(result1, [Row(key="key1", result="received:hello@10000")])

        # Process event at t=11000 for key2 - registers timer at t=16000.
        result2 = tester.test("key2", [Row(event_time_ms=11000, value="world")])
        self.assertEqual(result2, [Row(key="key2", result="received:world@11000")])

        # Set watermark to 15000 - key1's timer should fire.
        expired1 = tester.setWatermark(15000)
        self.assertEqual(len(expired1), 1)
        self.assertEqual(expired1[0], Row(key="key1", result="session-expired@watermark=15000"))

        # Set watermark to 16000 - key2's timer should fire.
        expired2 = tester.setWatermark(16000)
        self.assertEqual(len(expired2), 1)
        self.assertEqual(expired2[0], Row(key="key2", result="session-expired@watermark=16000"))

        # Verify state is cleared.
        self.assertIsNone(tester.peekValueState("lastEventTime", "key1"))
        self.assertIsNone(tester.peekValueState("lastEventTime", "key2"))

    def test_late_event_filtering(self):
        """Test that TwsTester filters late events in EventTime mode."""
        processor = EventTimeCountProcessorFactory().row()

        def event_time_extractor(row: Row) -> int:
            return row.event_time_ms

        tester = TwsTester(
            processor,
            timeMode="EventTime",
            eventTimeExtractor=event_time_extractor,
        )

        # Initially watermark is 0, so all events should be processed.
        result1 = tester.test(
            "key1",
            [
                Row(event_time_ms=5000, value="a"),
                Row(event_time_ms=6000, value="b"),
                Row(event_time_ms=7000, value="c"),
            ],
        )
        self.assertEqual(result1, [Row(key="key1", count=3)])
        self.assertEqual(tester.peekValueState("count", "key1"), (3,))

        # Set watermark to 6000.
        tester.setWatermark(6000)

        # Now watermark is 6000. Send events with mixed event times:
        # - (4000, "late1") -> event time 4000 <= watermark 6000, should be FILTERED.
        # - (5000, "late2") -> event time 5000 <= watermark 6000, should be FILTERED.
        # - (6000, "late3") -> event time 6000 <= watermark 6000, should be FILTERED.
        # - (8000, "ontime") -> event time 8000 > watermark 6000, should be PROCESSED.
        result2 = tester.test(
            "key1",
            [
                Row(event_time_ms=4000, value="late1"),
                Row(event_time_ms=5000, value="late2"),
                Row(event_time_ms=6000, value="late3"),
                Row(event_time_ms=8000, value="ontime"),
            ],
        )
        # Only 1 event should be processed (the on-time one)
        self.assertEqual(result2, [Row(key="key1", count=4)])  # 3 + 1 = 4
        self.assertEqual(tester.peekValueState("count", "key1"), (4,))

    def test_registers_on_time_events(self):
        "TwsTester with eventTimeExtractor should allow timer registration from on-time events"
        # With eventTimeExtractor, late events are filtered BEFORE reaching the processor,
        # so timers can only be registered from on-time events (no exception).
        tester = TwsTester(
            EventTimeSessionProcessorFactory().row(),
            timeMode="EventTime",
            eventTimeExtractor=lambda row: row.event_time_ms,
        )

        # Set watermark to 20 seconds.
        tester.setWatermark(20000)

        # Process a mix of late and on-time events.
        # Late event (10000) is filtered out, on-time event (25000) is processed.
        # Timer registered at 25000 + 5000 = 30000, which is > watermark, so no exception.
        result = tester.test(
            "key1",
            [Row(event_time_ms=10000, value="late"), Row(event_time_ms=25000, value="ontime")],
        )
        # Only the on-time event produces output.
        self.assertEqual(result, [Row(key="key1", result="received:ontime@25000")])

    def test_timer_registration_in_none_mode_raises_error(self):
        """Test that timer operations raise error in TimeMode.None."""
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor, timeMode="None")

        with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
            tester.handle.registerTimer(12345)

    def test_set_processing_time_in_wrong_mode_raises_error(self):
        """Test that setProcessingTime raises error in non-ProcessingTime mode."""
        processor = RunningCountStatefulProcessorFactory().row()

        # Test in None mode.
        tester_none = TwsTester(processor, timeMode="None")
        with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
            tester_none.setProcessingTime(1000)

    def test_set_watermark_in_wrong_mode_raises_error(self):
        """Test that setWatermark raises error in non-EventTime mode."""
        processor = RunningCountStatefulProcessorFactory().row()

        # Test in None mode.
        tester_none = TwsTester(processor, timeMode="None")
        with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
            tester_none.setWatermark(1000)

        # Test in ProcessingTime mode.
        tester_pt = TwsTester(processor, timeMode="ProcessingTime")
        with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
            tester_pt.setWatermark(1000)

    def test_eventtime_mode_requires_extractor(self):
        """Test that EventTime mode requires eventTimeExtractor."""
        processor = RunningCountStatefulProcessorFactory().row()

        with self.assertRaisesRegex(PySparkValueError, "ARGUMENT_REQUIRED"):
            TwsTester(processor, timeMode="EventTime")

    def test_processing_time_must_move_forward(self):
        """Test that setProcessingTime requires time to move forward."""
        processor = RunningCountStatefulProcessorFactory().row()
        tester = TwsTester(processor, timeMode="ProcessingTime")

        tester.setProcessingTime(1000)

        # Same time should raise error.
        with self.assertRaises(IllegalArgumentException):
            tester.setProcessingTime(1000)

        # Earlier time should raise error.
        with self.assertRaises(IllegalArgumentException):
            tester.setProcessingTime(500)

    def test_watermark_must_move_forward(self):
        """Test that setWatermark requires watermark to move forward."""
        processor = RunningCountStatefulProcessorFactory().row()

        def event_time_extractor(row: Row) -> int:
            return row.event_time_ms

        tester = TwsTester(
            processor,
            timeMode="EventTime",
            eventTimeExtractor=event_time_extractor,
        )

        tester.setWatermark(1000)

        # Same watermark should raise error.
        with self.assertRaises(IllegalArgumentException):
            tester.setWatermark(1000)

        # Earlier watermark should raise error.
        with self.assertRaises(IllegalArgumentException):
            tester.setWatermark(500)

    def test_timer_cannot_be_registered_in_past(self):
        """Test that timer cannot be registered at or before current watermark in EventTime mode."""
        processor = EventTimeSessionProcessorFactory().row()

        def event_time_extractor(row: Row) -> int:
            return row.event_time_ms

        tester = TwsTester(
            processor,
            timeMode="EventTime",
            eventTimeExtractor=event_time_extractor,
        )

        # Set watermark to 10000.
        tester.setWatermark(10000)

        # Try to register a timer at or before watermark - should raise error.
        tester.handle.setGroupingKey("key1")
        with self.assertRaises(IllegalArgumentException):
            tester.handle.registerTimer(10000)  # exactly at watermark

        with self.assertRaises(IllegalArgumentException):
            tester.handle.registerTimer(5000)  # before watermark


@unittest.skipIf(
    not have_pandas or not have_pyarrow or os.environ.get("PYTHON_GIL", "?") == "0",
    pandas_requirement_message or pyarrow_requirement_message or "Not supported in no-GIL mode",
)
class TwsTesterFuzzTests(ReusedSQLTestCase):
    @classmethod
    def conf(cls):
        cfg = SparkConf()
        cfg.set("spark.sql.shuffle.partitions", "1")
        cfg.set(
            "spark.sql.streaming.stateStore.providerClass",
            "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider",
        )
        return cfg

    # Helper that runs real streaming query with TWS and compares results with TwsTester output.
    # Supports both row mode (use_pandas=False) and Pandas mode (use_pandas=True).
    # Supports multiple batches.
    # Assumes that the first column is the key column.
    # Output rows are checked ignoring order.
    # Correct only if results do not depend on reordering rows within batch.
    def _check_tws_tester(
        self,
        processor: StatefulProcessor,
        batches: list[list[tuple]],
        input_columns: list[str],
        input_types: list[str],
        output_schema: StructType,
        use_pandas: bool = False,
    ):
        num_batches = len(batches)
        key_column: str = input_columns[0]  # Assume first column is the key column.
        value_columns: list[str] = input_columns[1:]  # All columns except the key column.

        num_input_columns: int = len(input_columns)
        assert len(input_types) == num_input_columns

        # Use TwsTester to compute expected results.
        tester = TwsTester(processor)

        def _run_tester(batch: list[tuple]) -> list[Row]:
            # Group batch by key.
            grouped = {}
            for row in batch:
                key = row[0]
                if key not in grouped:
                    grouped[key] = []
                grouped[key].append(row[1:])

            results: list[Row] = []
            for key, values in grouped.items():
                if use_pandas:
                    input_to_tester = pd.DataFrame(values, columns=value_columns)
                    result_df = tester.testInPandas(key, input_to_tester)
                    results.extend(self.spark.createDataFrame(result_df).collect())
                else:
                    input_to_tester: list[Row] = [
                        Row(**dict(zip(value_columns, v))) for v in values
                    ]
                    results.extend(tester.test(key, input_to_tester))
            return results

        expected_results: list[list[Row]] = [_run_tester(batch) for batch in batches]

        # Create streaming DataFrame, it will read input from files.
        input_path = tempfile.mkdtemp()
        df = self.spark.readStream.format("text").load(input_path)
        df_split = df.withColumn("split_values", split(df["value"], ","))
        df_final = df_split.select(
            *[
                df_split.split_values.getItem(i).alias(col_name).cast(input_types[i])
                for i, col_name in enumerate(input_columns)
            ]
        )
        self.assertTrue(df_final.isStreaming)

        # Apply TransformWithState.
        df_grouped = df_final.groupBy(key_column)
        if use_pandas:
            tws_df = df_grouped.transformWithStateInPandas(
                statefulProcessor=processor,
                outputStructType=output_schema,
                outputMode="Update",
                timeMode="None",
            )
        else:
            tws_df = df_grouped.transformWithState(
                statefulProcessor=processor,
                outputStructType=output_schema,
                outputMode="Update",
                timeMode="None",
            )

        # Start streaming query collecting TWS results.
        actual_results: list[list[Row]] = []

        def _collect_results(batch_df: DataFrame, batch_id):
            rows: list[Row] = batch_df.collect()
            actual_results.append(rows)

        query: StreamingQuery = (
            tws_df.writeStream.queryName("test_tmp_query")
            .foreachBatch(_collect_results)
            .outputMode("update")
            .start()
        )

        # Process each input batch.
        def _write_batch(batch_id: int, batch: list[tuple]):
            with open(input_path + f"/test{batch_id}.txt", "w") as f:
                for row in batch:
                    assert len(row) == num_input_columns
                    f.write(",".join(str(value) for value in row) + "\n")

        try:
            for batch_id, batch in enumerate(batches):
                _write_batch(batch_id, batch)
                query.processAllAvailable()
            query.awaitTermination(1)
            self.assertTrue(query.exception() is None)
        finally:
            query.stop()

        # Assert actual results are equal to expected results.
        self.assertEqual(len(actual_results), num_batches)
        self.assertEqual(len(expected_results), num_batches)
        for batch_id in range(num_batches):
            # This checks that collections are the same up to order, not just compares lengths.
            self.assertCountEqual(actual_results[batch_id], expected_results[batch_id])

    def test_fuzz_running_count(self):
        output_schema = StructType(
            [
                StructField("key", StringType(), True),
                StructField("count", IntegerType(), True),
            ]
        )
        batches: list[list[tuple[str, str]]] = []
        for batch_id in range(5):
            batch: list[tuple[str, str]] = []
            for _ in range(200):
                key: str = f"key{random.randint(0, 9)}"
                value: str = "".join(random.choices("abcdefghijklmnopqrstuvwxyz", k=5))
                batch.append((key, value))
            batches.append(batch)
        input_columns: list[str] = ["key", "value"]
        input_types: list[str] = ["string", "string"]

        for use_pandas in [False, True]:
            self._check_tws_tester(
                RunningCountStatefulProcessorFactory().get(use_pandas=use_pandas),
                batches,
                input_columns,
                input_types,
                output_schema,
                use_pandas=use_pandas,
            )

    def test_fuzz_topk(self):
        output_schema = StructType(
            [
                StructField("key", StringType(), True),
                StructField("score", DoubleType(), True),
            ]
        )
        batches: list[list[tuple[str, float]]] = []
        for batch_id in range(5):
            batch: list[tuple[str, float]] = []
            for _ in range(200):
                key: str = f"key{random.randint(0, 9)}"
                score: float = random.randint(0, 1000000) / 1000
                batch.append((key, score))
            batches.append(batch)
        input_columns: list[str] = ["key", "score"]
        input_types: list[str] = ["string", "double"]

        for use_pandas in [False, True]:
            self._check_tws_tester(
                TopKProcessorFactory(k=10).get(use_pandas=use_pandas),
                batches,
                input_columns,
                input_types,
                output_schema,
                use_pandas=use_pandas,
            )

    def test_fuzz_word_frequency(self):
        output_schema = StructType(
            [
                StructField("key", StringType(), True),
                StructField("word", StringType(), True),
                StructField("count", IntegerType(), True),
            ]
        )
        batches: list[list[tuple[str, str]]] = []
        for batch_id in range(5):
            batch: list[tuple[str, str]] = []
            for _ in range(200):
                key = f"key{random.randint(0, 9)}"
                word = "".join(random.choices("abcdefghijklmnopqrstuvwxyz", k=1))
                batch.append((key, word))
            batches.append(batch)
        input_columns: list[str] = ["key", "word"]
        input_types: list[str] = ["string", "string"]

        for use_pandas in [False, True]:
            self._check_tws_tester(
                WordFrequencyProcessorFactory().get(use_pandas=use_pandas),
                batches,
                input_columns,
                input_types,
                output_schema,
                use_pandas=use_pandas,
            )


if __name__ == "__main__":
    from pyspark.testing import main

    main()
