#
# 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 tempfile
import unittest
from typing import cast

from pyspark import SparkConf
from pyspark.sql.functions import split
from pyspark.sql.streaming.state import GroupStateTimeout
from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
    SimpleStatefulProcessorWithInitialStateFactory,
    StatefulProcessorCompositeTypeFactory,
)
from pyspark.sql.types import LongType, StringType, StructType, StructField
from pyspark.testing.sqlutils import ReusedSQLTestCase
from pyspark.testing.utils import (
    have_pandas,
    have_pyarrow,
    pandas_requirement_message,
    pyarrow_requirement_message,
)

if have_pandas:
    import pandas as pd

if have_pyarrow:
    import pyarrow as pa  # noqa: F401


class OfflineStateRepartitionTestUtils:
    """Utility class for repartition tests."""

    @staticmethod
    def run_repartition_test(
        spark,
        num_shuffle_partitions,
        create_streaming_df,
        output_mode,
        batch1_data,
        batch2_data,
        batch3_data,
        verify_initial,
        verify_after_increase,
        verify_after_decrease,
    ):
        """
        Common helper to run repartition tests with different streaming operators.

        Steps:
        1. Run streaming query to generate initial state
        2. Repartition to more partitions
        3. Add new data and restart query, verify state preserved
        4. Repartition to fewer partitions
        5. Add new data and restart query, verify state preserved

        Parameters
        ----------
        spark : SparkSession
            The active Spark session.
        num_shuffle_partitions : int
            The initial number of shuffle partitions.
        create_streaming_df : callable
            Function(df) -> DataFrame that applies the streaming transformation.
        output_mode : str
            Output mode for the streaming query ("update" or "append").
        batch1_data, batch2_data, batch3_data : str
            Data to write for each batch (newline-separated values).
        verify_initial, verify_after_increase, verify_after_decrease : callable
            Functions(collected_results) to verify results at each stage.
        """
        with (
            tempfile.TemporaryDirectory() as input_dir,
            tempfile.TemporaryDirectory() as checkpoint_dir,
        ):
            collected_results = []

            def collect_batch(batch_df, batch_id):
                rows = batch_df.collect()
                collected_results.extend(rows)

            def run_streaming_query():
                df = spark.readStream.format("text").load(input_dir)
                transformed_df = create_streaming_df(df)
                query = (
                    transformed_df.writeStream.foreachBatch(collect_batch)
                    .option("checkpointLocation", checkpoint_dir)
                    .outputMode(output_mode)
                    .start()
                )
                query.processAllAvailable()
                query.stop()

            # Step 1: Write initial data and run streaming query
            with open(os.path.join(input_dir, "batch1.txt"), "w") as f:
                f.write(batch1_data)

            run_streaming_query()
            verify_initial(collected_results)

            # Step 2: Repartition to more partitions
            spark._streamingCheckpointManager.repartition(
                checkpoint_dir, num_shuffle_partitions * 2
            )

            # Step 3: Add more data and restart query
            with open(os.path.join(input_dir, "batch2.txt"), "w") as f:
                f.write(batch2_data)

            collected_results.clear()
            run_streaming_query()
            verify_after_increase(collected_results)

            # Step 4: Repartition to fewer partitions
            spark._streamingCheckpointManager.repartition(
                checkpoint_dir, num_shuffle_partitions - 1
            )

            # Step 5: Add more data and restart query
            with open(os.path.join(input_dir, "batch3.txt"), "w") as f:
                f.write(batch3_data)

            collected_results.clear()
            run_streaming_query()
            verify_after_decrease(collected_results)


class StreamingOfflineStateRepartitionTests(ReusedSQLTestCase):
    """
    Test suite for Offline state repartitioning.
    """

    NUM_SHUFFLE_PARTITIONS = 3

    @classmethod
    def conf(cls):
        cfg = SparkConf()
        cfg.set("spark.sql.shuffle.partitions", str(cls.NUM_SHUFFLE_PARTITIONS))
        cfg.set(
            "spark.sql.streaming.stateStore.providerClass",
            "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider",
        )
        return cfg

    def test_fail_if_empty_checkpoint_directory(self):
        """Test that repartition fails if checkpoint directory is empty."""
        with tempfile.TemporaryDirectory() as checkpoint_dir:
            with self.assertRaisesRegex(
                Exception, "STATE_REPARTITION_INVALID_CHECKPOINT.NO_COMMITTED_BATCH"
            ):
                self.spark._streamingCheckpointManager.repartition(checkpoint_dir, 5)

    def test_fail_if_no_batch_found_in_checkpoint_directory(self):
        """Test that repartition fails if no batch found in checkpoint directory."""
        with tempfile.TemporaryDirectory() as checkpoint_dir:
            # Write commit log but no offset log
            commits_dir = os.path.join(checkpoint_dir, "commits")
            os.makedirs(commits_dir)
            # Create a minimal commit file for batch 0
            with open(os.path.join(commits_dir, "0"), "w") as f:
                f.write("v1\n{}")

            with self.assertRaisesRegex(
                Exception, "STATE_REPARTITION_INVALID_CHECKPOINT.NO_BATCH_FOUND"
            ):
                self.spark._streamingCheckpointManager.repartition(checkpoint_dir, 5)

    def test_fail_if_repartition_parameter_is_invalid(self):
        """Test that repartition fails with invalid parameters."""
        # Test null checkpoint location
        with self.assertRaisesRegex(Exception, "STATE_REPARTITION_INVALID_PARAMETER.IS_NULL"):
            self.spark._streamingCheckpointManager.repartition(None, 5)

        # Test empty checkpoint location
        with self.assertRaisesRegex(Exception, "STATE_REPARTITION_INVALID_PARAMETER.IS_EMPTY"):
            self.spark._streamingCheckpointManager.repartition("", 5)

        # Test numPartitions <= 0
        with self.assertRaisesRegex(
            Exception, "STATE_REPARTITION_INVALID_PARAMETER.IS_NOT_GREATER_THAN_ZERO"
        ):
            self.spark._streamingCheckpointManager.repartition("test", 0)

    @unittest.skipIf(
        not have_pandas or not have_pyarrow,
        pandas_requirement_message or pyarrow_requirement_message,
    )
    def test_repartition_with_apply_in_pandas_with_state(self):
        """Test repartition for a streaming query using applyInPandasWithState."""

        # Define the stateful function that tracks count per key
        def stateful_count_func(key, pdf_iter, state):
            existing_count = state.getOption
            count = 0 if existing_count is None else existing_count[0]

            new_count = 0
            for pdf in pdf_iter:
                new_count += len(pdf)

            total_count = count + new_count
            state.update((total_count,))
            yield pd.DataFrame({"key": [key[0]], "count": [total_count]})

        output_schema = StructType(
            [StructField("key", StringType()), StructField("count", LongType())]
        )
        state_schema = StructType([StructField("count", LongType())])

        def create_streaming_df(df):
            return df.groupBy(df["value"]).applyInPandasWithState(
                stateful_count_func,
                output_schema,
                state_schema,
                "Update",
                GroupStateTimeout.NoTimeout,
            )

        def verify_initial(results):
            counts = {row.key: row["count"] for row in results}
            self.assertEqual(counts.get("a"), 2)
            self.assertEqual(counts.get("b"), 1)
            self.assertEqual(counts.get("c"), 1)

        def verify_after_increase(results):
            counts = {row.key: row["count"] for row in results}
            self.assertEqual(counts.get("a"), 3, "State for 'a': 2 + 1 = 3")
            self.assertEqual(counts.get("b"), 2, "State for 'b': 1 + 1 = 2")
            self.assertEqual(counts.get("d"), 1, "New key 'd' should have count 1")

        def verify_after_decrease(results):
            counts = {row.key: row["count"] for row in results}
            self.assertEqual(counts.get("a"), 4, "State for 'a': 3 + 1 = 4")
            self.assertEqual(counts.get("c"), 2, "State for 'c': 1 + 1 = 2")
            self.assertEqual(counts.get("e"), 1, "New key 'e' should have count 1")

        OfflineStateRepartitionTestUtils.run_repartition_test(
            spark=self.spark,
            num_shuffle_partitions=self.NUM_SHUFFLE_PARTITIONS,
            create_streaming_df=create_streaming_df,
            output_mode="update",
            batch1_data="a\nb\na\nc\n",  # a:2, b:1, c:1
            batch2_data="a\nb\nd\n",  # a:+1, b:+1, d:1 (new)
            batch3_data="a\nc\ne\n",  # a:+1, c:+1, e:1 (new)
            verify_initial=verify_initial,
            verify_after_increase=verify_after_increase,
            verify_after_decrease=verify_after_decrease,
        )

    def test_repartition_with_streaming_aggregation(self):
        """Test repartition for a streaming aggregation query (groupBy + count)."""

        def create_streaming_df(df):
            return df.groupBy("value").count()

        def verify_initial(results):
            counts = {row.value: row["count"] for row in results}
            self.assertEqual(counts.get("a"), 2)
            self.assertEqual(counts.get("b"), 1)
            self.assertEqual(counts.get("c"), 1)

        def verify_after_increase(results):
            counts = {row.value: row["count"] for row in results}
            self.assertEqual(counts.get("a"), 3, "State for 'a': 2 + 1 = 3")
            self.assertEqual(counts.get("b"), 2, "State for 'b': 1 + 1 = 2")
            self.assertEqual(counts.get("d"), 1, "New key 'd' should have count 1")

        def verify_after_decrease(results):
            counts = {row.value: row["count"] for row in results}
            self.assertEqual(counts.get("a"), 4, "State for 'a': 3 + 1 = 4")
            self.assertEqual(counts.get("c"), 2, "State for 'c': 1 + 1 = 2")
            self.assertEqual(counts.get("e"), 1, "New key 'e' should have count 1")

        OfflineStateRepartitionTestUtils.run_repartition_test(
            spark=self.spark,
            num_shuffle_partitions=self.NUM_SHUFFLE_PARTITIONS,
            create_streaming_df=create_streaming_df,
            output_mode="update",
            batch1_data="a\nb\na\nc\n",  # a:2, b:1, c:1
            batch2_data="a\nb\nd\n",  # a:+1, b:+1, d:1 (new)
            batch3_data="a\nc\ne\n",  # a:+1, c:+1, e:1 (new)
            verify_initial=verify_initial,
            verify_after_increase=verify_after_increase,
            verify_after_decrease=verify_after_decrease,
        )

    def test_repartition_with_streaming_dedup(self):
        """Test repartition for a streaming deduplication query (dropDuplicates)."""

        def create_streaming_df(df):
            return df.dropDuplicates(["value"])

        def verify_initial(results):
            values = {row.value for row in results}
            self.assertEqual(values, {"a", "b", "c"})

        def verify_after_increase(results):
            values = {row.value for row in results}
            self.assertEqual(values, {"d", "e"}, "Only new keys after repartition")

        def verify_after_decrease(results):
            values = {row.value for row in results}
            self.assertEqual(values, {"f", "g"}, "Only new keys after repartition")

        OfflineStateRepartitionTestUtils.run_repartition_test(
            spark=self.spark,
            num_shuffle_partitions=self.NUM_SHUFFLE_PARTITIONS,
            create_streaming_df=create_streaming_df,
            output_mode="append",
            batch1_data="a\nb\na\nc\n",  # unique: a, b, c
            batch2_data="a\nb\nd\ne\n",  # a, b duplicates; d, e new
            batch3_data="a\nc\nf\ng\n",  # a, c duplicates; f, g new
            verify_initial=verify_initial,
            verify_after_increase=verify_after_increase,
            verify_after_decrease=verify_after_decrease,
        )

    def _run_tws_repartition_test(self, is_pandas):
        """Helper method to run repartition test with a given processor factory method"""

        def create_streaming_df(df):
            # Parse text input format "id,temperature" into structured columns
            split_df = split(df["value"], ",")
            parsed_df = df.select(
                split_df.getItem(0).alias("id"),
                split_df.getItem(1).cast("integer").alias("temperature"),
            )

            output_schema = StructType(
                [
                    StructField("id", StringType(), True),
                    StructField("value", StringType(), True),
                ]
            )
            processor = (
                SimpleStatefulProcessorWithInitialStateFactory().pandas()
                if is_pandas
                else SimpleStatefulProcessorWithInitialStateFactory().row()
            )
            return (
                parsed_df.groupBy("id").transformWithStateInPandas(
                    statefulProcessor=processor,
                    outputStructType=output_schema,
                    outputMode="Update",
                    timeMode="None",
                    initialState=None,
                )
                if is_pandas
                else parsed_df.groupBy("id").transformWithState(
                    statefulProcessor=processor,
                    outputStructType=output_schema,
                    outputMode="Update",
                    timeMode="None",
                    initialState=None,
                )
            )

        def verify_initial(results):
            # SimpleStatefulProcessorWithInitialState accumulates temperature values
            values = {row.id: row.value for row in results}
            self.assertEqual(values.get("a"), "270", "a: 120 + 150 = 270")
            self.assertEqual(values.get("b"), "50", "b: 50")
            self.assertEqual(values.get("c"), "30", "c: 30")

        def verify_after_increase(results):
            # After repartition, state should be preserved and accumulated
            values = {row.id: row.value for row in results}
            self.assertEqual(values.get("a"), "371", "State for 'a': 270 + 101 = 371")
            self.assertEqual(values.get("b"), "152", "State for 'b': 50 + 102 = 152")
            self.assertEqual(values.get("d"), "103", "New key 'd' should have value 103")

        def verify_after_decrease(results):
            # After repartition, state should still be preserved
            values = {row.id: row.value for row in results}
            self.assertEqual(values.get("a"), "475", "State for 'a': 371 + 104 = 475")
            self.assertEqual(values.get("c"), "135", "State for 'c': 30 + 105 = 135")
            self.assertEqual(values.get("e"), "106", "New key 'e' should have value 106")

        OfflineStateRepartitionTestUtils.run_repartition_test(
            spark=self.spark,
            num_shuffle_partitions=self.NUM_SHUFFLE_PARTITIONS,
            create_streaming_df=create_streaming_df,
            output_mode="update",
            batch1_data="a,120\na,150\nb,50\nc,30\n",  # a:270, b:50, c:30
            batch2_data="a,101\nb,102\nd,103\n",  # a:371, b:152, d:103 (new)
            batch3_data="a,104\nc,105\ne,106\n",  # a:475, c:135, e:106 (new)
            verify_initial=verify_initial,
            verify_after_increase=verify_after_increase,
            verify_after_decrease=verify_after_decrease,
        )

    @unittest.skipIf(
        not have_pyarrow or os.environ.get("PYTHON_GIL", "?") == "0",
        pyarrow_requirement_message or "Not supported in no-GIL mode",
    )
    def test_repartition_with_streaming_tws(self):
        """Test repartition for streaming transformWithState"""

        self._run_tws_repartition_test(is_pandas=False)

    @unittest.skipIf(
        not have_pandas or not have_pyarrow or os.environ.get("PYTHON_GIL", "?") == "0",
        cast(
            str,
            pandas_requirement_message
            or pyarrow_requirement_message
            or "Not supported in no-GIL mode",
        ),
    )
    def test_repartition_with_streaming_tws_in_pandas(self):
        """Test repartition for streaming transformWithStateInPandas."""

        self._run_tws_repartition_test(is_pandas=True)

    def _run_repartition_with_streaming_tws_multiple_state_vars_test(self, is_pandas):
        """Test repartition with processor using multiple state variable types (value + list + map)."""

        def create_streaming_df(df):
            # Parse text input format "id,temperature" into structured columns
            split_df = split(df["value"], ",")
            parsed_df = df.select(
                split_df.getItem(0).alias("id"),
                split_df.getItem(1).cast("integer").alias("temperature"),
            )

            output_schema = StructType(
                [
                    StructField("id", StringType(), True),
                    StructField("value_arr", StringType(), True),
                    StructField("list_state_arr", StringType(), True),
                    StructField("map_state_arr", StringType(), True),
                    StructField("nested_map_state_arr", StringType(), True),
                ]
            )
            processor = (
                StatefulProcessorCompositeTypeFactory().pandas()
                if is_pandas
                else StatefulProcessorCompositeTypeFactory().row()
            )
            return (
                parsed_df.groupBy("id").transformWithStateInPandas(
                    statefulProcessor=processor,
                    outputStructType=output_schema,
                    outputMode="Update",
                    timeMode="None",
                    initialState=None,
                )
                if is_pandas
                else parsed_df.groupBy("id").transformWithState(
                    statefulProcessor=processor,
                    outputStructType=output_schema,
                    outputMode="Update",
                    timeMode="None",
                    initialState=None,
                )
            )

        def verify_initial(results):
            rows = {row.id: row for row in results}
            # StatefulProcessorCompositeType initializes to "0" on first batch
            # (this is how the processor is designed - see line 1853 in helper file)
            self.assertEqual(rows["a"].value_arr, "0")
            self.assertEqual(rows["a"].list_state_arr, "0")
            # Map state initialized with default ATTRIBUTES_MAP and CONFS_MAP
            self.assertEqual(rows["a"].map_state_arr, '{"key1": [1], "key2": [10]}')
            self.assertEqual(rows["a"].nested_map_state_arr, '{"e1": {"e2": 5, "e3": 10}}')

        def verify_after_increase(results):
            # After repartition, state should be preserved
            rows = {row.id: row for row in results}
            # Key 'a': value state accumulated
            self.assertEqual(rows["a"].value_arr, "100")
            self.assertEqual(rows["a"].list_state_arr, "0,100")
            # Map state updated with key "a" and temperature 100
            self.assertEqual(rows["a"].map_state_arr, '{"a": [100], "key1": [1], "key2": [10]}')
            self.assertEqual(
                rows["a"].nested_map_state_arr, '{"e1": {"a": 100, "e2": 5, "e3": 10}}'
            )
            # New key 'd' - first batch for this key, outputs "0"
            self.assertEqual(rows["d"].value_arr, "0")
            self.assertEqual(rows["d"].list_state_arr, "0")
            # Map state for 'd' initialized with defaults
            self.assertEqual(rows["d"].map_state_arr, '{"key1": [1], "key2": [10]}')
            self.assertEqual(rows["d"].nested_map_state_arr, '{"e1": {"e2": 5, "e3": 10}}')

        def verify_after_decrease(results):
            # After another repartition, state should still be preserved
            rows = {row.id: row for row in results}
            # Key 'd' gets 102, state was [0], so accumulated: [0+102] = [102]
            self.assertEqual(rows["d"].value_arr, "102")
            self.assertEqual(rows["d"].list_state_arr, "0,102")
            # Map state for 'd' updated with temperature 102
            self.assertEqual(rows["d"].map_state_arr, '{"d": [102], "key1": [1], "key2": [10]}')
            self.assertEqual(
                rows["d"].nested_map_state_arr, '{"e1": {"d": 102, "e2": 5, "e3": 10}}'
            )
            # Key 'a' gets 101, state was [100], so accumulated: [100+101] = [201]
            self.assertEqual(rows["a"].value_arr, "201")
            self.assertEqual(rows["a"].list_state_arr, "0,100,101")
            # Map state for 'a' updated with temperature 101 (replaces previous value)
            self.assertEqual(rows["a"].map_state_arr, '{"a": [101], "key1": [1], "key2": [10]}')
            self.assertEqual(
                rows["a"].nested_map_state_arr, '{"e1": {"a": 101, "e2": 5, "e3": 10}}'
            )

        OfflineStateRepartitionTestUtils.run_repartition_test(
            spark=self.spark,
            num_shuffle_partitions=self.NUM_SHUFFLE_PARTITIONS,
            create_streaming_df=create_streaming_df,
            output_mode="update",
            batch1_data="a,100\n",  # a:0
            batch2_data="a,100\nd,100\n",  # a:100, d:0 (new)
            batch3_data="d,102\na,101\n",  # d:102, a:201 (new)
            verify_initial=verify_initial,
            verify_after_increase=verify_after_increase,
            verify_after_decrease=verify_after_decrease,
        )

    @unittest.skipIf(
        not have_pyarrow or os.environ.get("PYTHON_GIL", "?") == "0",
        pyarrow_requirement_message or "Not supported in no-GIL mode",
    )
    def test_repartition_with_streaming_tws_multiple_state_vars(self):
        """Test repartition with transformWithState using multiple state variable types (value + list + map)."""

        self._run_repartition_with_streaming_tws_multiple_state_vars_test(is_pandas=False)

    @unittest.skipIf(
        not have_pandas or not have_pyarrow or os.environ.get("PYTHON_GIL", "?") == "0",
        cast(
            str,
            pandas_requirement_message
            or pyarrow_requirement_message
            or "Not supported in no-GIL mode",
        ),
    )
    def test_repartition_with_streaming_tws_in_pandas_multiple_state_vars(self):
        """Test repartition with transformWithStateInPandas using multiple state variable types (value + list + map)."""

        self._run_repartition_with_streaming_tws_multiple_state_vars_test(is_pandas=True)


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

    main()
