/*
 * 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.
 */
package org.apache.spark.sql.execution.streaming.state

import java.io.File
import java.sql.Timestamp
import java.time.Duration

import org.apache.spark.sql.Row
import org.apache.spark.sql.execution.datasources.v2.state.{StateDataSourceTestBase, StateSourceOptions, StreamStreamJoinTestUtils}
import org.apache.spark.sql.execution.streaming.operators.stateful.{StatefulOperatorsUtils, StreamingAggregationStateManager}
import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.timers.TimerStateUtils
import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream, StreamingQueryCheckpointMetadata}
import org.apache.spark.sql.execution.streaming.utils.StreamingUtils
import org.apache.spark.sql.functions.{col, timestamp_seconds}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.streaming.{InputEvent, ListStateTTLProcessor, MapInputEvent, MapStateTTLProcessor, OutputMode, RunningCountStatefulProcessorWithProcTimeTimer, TimeMode, Trigger, TTLConfig, ValueStateTTLProcessor}
import org.apache.spark.sql.streaming.util.{StreamManualClock, TTLProcessorUtils}
import org.apache.spark.sql.streaming.util.{EventTimeTimerProcessor, MultiStateVarProcessor, MultiStateVarProcessorTestUtils, TimerTestUtils}
import org.apache.spark.sql.types.StructType

/**
 * Test suite for StatePartitionAllColumnFamiliesWriter.
 * Tests the writer's ability to correctly write raw bytes read from
 * StatePartitionAllColumnFamiliesReader to a state store without loading previous versions.
 */
class StatePartitionAllColumnFamiliesWriterSuite extends StateDataSourceTestBase {
  import testImplicits._

  override def beforeAll(): Unit = {
    super.beforeAll()
    spark.conf.set(SQLConf.STATE_STORE_PROVIDER_CLASS.key,
      classOf[RocksDBStateStoreProvider].getName)
    spark.conf.set(SQLConf.SHUFFLE_PARTITIONS.key, "2")
  }

  /**
   * Common helper method to perform round-trip test: read state bytes from source,
   * write to target, and verify target matches source.
   *
   * @param sourceDir Source checkpoint directory
   * @param targetDir Target checkpoint directory
   * @param storeToColumnFamilies Optional store name to its column families
   * @param storeToColumnFamilyToSelectExprs Map store name to per column family custom selectExprs
   * @param storeToColumnFamilyToStateSourceOptions Map store name to per column family
   *                                                state source options
   */
  private def performRoundTripTest(
      sourceDir: String,
      targetDir: String,
      storeToColumnFamilies: Map[String, List[String]] =
        Map(StateStoreId.DEFAULT_STORE_NAME -> List(StateStore.DEFAULT_COL_FAMILY_NAME)),
      storeToColumnFamilyToSelectExprs: Map[String, Map[String, Seq[String]]] = Map.empty,
      storeToColumnFamilyToStateSourceOptions: Map[String, Map[String, Map[String, String]]] =
        Map.empty,
      operatorName: String): Unit = {
    val hadoopConf = spark.sessionState.newHadoopConf()
    val sourceCpLocation = StreamingUtils.resolvedCheckpointLocation(
      hadoopConf, sourceDir)
    val sourceCheckpointMetadata = new StreamingQueryCheckpointMetadata(
      spark, sourceCpLocation)
    val readBatchId = sourceCheckpointMetadata.commitLog.getLatestBatchId().get

    val targetCpLocation = StreamingUtils.resolvedCheckpointLocation(
      hadoopConf, targetDir)
    val targetCheckpointMetadata = new StreamingQueryCheckpointMetadata(
      spark, targetCpLocation)
    // increase offsetCheckpoint
    val lastBatch = targetCheckpointMetadata.commitLog.getLatestBatchId().get
    val targetOffsetSeq = targetCheckpointMetadata.offsetLog.get(lastBatch).get
    val writeBatchId = lastBatch + 1
    targetCheckpointMetadata.offsetLog.add(writeBatchId, targetOffsetSeq)

    val rewriter = new StateRewriter(
      spark,
      readBatchId,
      writeBatchId,
      targetCpLocation,
      hadoopConf,
      readResolvedCheckpointLocation = Some(sourceCpLocation),
      transformFunc = None,
      writeCheckpointMetadata = Some(targetCheckpointMetadata)
    )
    val checkpointInfos = rewriter.run()

    // Commit to commitLog with checkpoint IDs
    val latestCommit = targetCheckpointMetadata.commitLog.get(lastBatch).get
    val commitMetadata = latestCommit.withStateUniqueIds(checkpointInfos)
    targetCheckpointMetadata.commitLog.add(writeBatchId, commitMetadata)
    val versionToCheck = writeBatchId + 1

    storeToColumnFamilies.foreach { case (storeName, columnFamilies) =>
      val storeNamePath = if (storeName == StateStoreId.DEFAULT_STORE_NAME) {
        "state/0/0"
      } else {
        s"state/0/0/$storeName"
      }
      assert(!checkpointFileExists(new File(targetDir, storeNamePath),
        versionToCheck, ".changelog"))
      assert(checkpointFileExists(new File(targetDir, storeNamePath), versionToCheck, ".zip"))

      // Validate by reading from both source and target using normal reader"
      // Default selectExprs for most column families
      val defaultSelectExprs = Seq("key", "value", "partition_id")

      columnFamilies
      // filtering out "default" for TWS operator because it doesn't contain any data
      .filter(cfName => !(cfName == StateStore.DEFAULT_COL_FAMILY_NAME &&
        StatefulOperatorsUtils.TRANSFORM_WITH_STATE_OP_NAMES.contains(operatorName)
      ))
      .foreach { cfName =>
        val selectExprs = storeToColumnFamilyToSelectExprs.getOrElse(storeName, Map.empty)
          .getOrElse(cfName, defaultSelectExprs)
        val readerOptions = storeToColumnFamilyToStateSourceOptions.getOrElse(storeName, Map.empty)
          .getOrElse(cfName, Map.empty)

        def readNormalData(dir: String): Array[Row] = {
          var reader = spark.read
            .format("statestore")
            .option(StateSourceOptions.PATH, dir)
            .option(StateSourceOptions.STORE_NAME, storeName)
          readerOptions.foreach { case (k, v) => reader = reader.option(k, v) }
          reader.load()
            .selectExpr(selectExprs: _*)
            .collect()
        }

        val sourceNormalData = readNormalData(sourceDir)
        val targetNormalData = readNormalData(targetDir)

        validateDataMatches(sourceNormalData, targetNormalData)
      }
    }
  }

  /**
   * Helper method to validate that source and target data match.
   */
  private def validateDataMatches(
      sourceNormalData: Array[Row],
      targetNormalData: Array[Row]): Unit = {
    assert(sourceNormalData.length == targetNormalData.length,
      s"Row count mismatch: source=${sourceNormalData.length}, " +
        s"target=${targetNormalData.length}")

    // Sort and compare row by row
    val sourceSorted = sourceNormalData.sortBy(_.toString)
    val targetSorted = targetNormalData.sortBy(_.toString)

    sourceSorted.zip(targetSorted).zipWithIndex.foreach {
      case ((sourceRow, targetRow), idx) =>
        assert(sourceRow == targetRow,
          s"Row mismatch at index $idx:\n" +
            s"  Source: $sourceRow\n" +
            s"  Target: $targetRow")
    }
  }

    /**
     * Checks if a changelog file for the specified version exists in the given directory.
     * A changelog file has the suffix ".changelog".
     *
     * @param dir Directory to search for changelog files
     * @param version The version to check for existence
     * @param suffix Either 'zip' or 'changelog'
     * @return true if a changelog file with the given version exists, false otherwise
     */
    private def checkpointFileExists(dir: File, version: Long, suffix: String): Boolean = {
      Option(dir.listFiles)
        .getOrElse(Array.empty)
        .map { file =>
          file
        }
        .filter { file =>
          file.getName.endsWith(suffix) && !file.getName.startsWith(".")
        }
        .exists { file =>
          val nameWithoutSuffix = file.getName.stripSuffix(suffix)
          val parts = nameWithoutSuffix.split("_")
          parts.headOption match {
            case Some(ver) if ver.forall(_.isDigit) => ver.toLong == version
            case _ => false
          }
        }
    }

  private def createColFamilyInfo(
       keySchema: StructType,
       valueSchema: StructType,
       keyStateEncoderSpec: KeyStateEncoderSpec,
       colFamilyName: String,
       useMultipleValuePerKey: Boolean = false): StatePartitionWriterColumnFamilyInfo = {
    StatePartitionWriterColumnFamilyInfo(
      schema = StateStoreColFamilySchema(
        colFamilyName,
        keySchemaId = 0,
        keySchema,
        valueSchemaId = 0,
        valueSchema,
        keyStateEncoderSpec = Some(keyStateEncoderSpec)
      ),
      useMultipleValuePerKey)
  }

  /**
   * Helper method to test SPARK-54420 read and write with different state format versions
   * for simple aggregation (single grouping key).
   * @param stateVersion The state format version (1 or 2)
   */
  private def testRoundTripForAggrStateVersion(stateVersion: Int): Unit = {
    withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> stateVersion.toString) {
      withTempDir { sourceDir =>
        withTempDir { targetDir =>
          // Step 1: Create state by running a streaming aggregation
          runLargeDataStreamingAggregationQuery(sourceDir.getAbsolutePath)
          val inputData: MemoryStream[Int] = MemoryStream[Int]
          val aggregated = getLargeDataStreamingAggregationQuery(inputData)

          // add dummy data to target source to test writer won't load previous store
          testStream(aggregated, OutputMode.Update)(
            StartStream(checkpointLocation = targetDir.getAbsolutePath),
            // batch 0
            AddData(inputData, 0 until 2: _*),
            CheckLastBatch(
              (0, 1, 0, 0, 0), // 0
              (1, 1, 1, 1, 1) // 1
            ),
            // batch 1
            AddData(inputData, 0 until 2: _*),
            CheckLastBatch(
              (0, 2, 0, 0, 0), // 0
              (1, 2, 2, 1, 1) // 1
            )
          )

          // Perform round-trip test using common helper
          performRoundTripTest(
            sourceDir.getAbsolutePath,
            targetDir.getAbsolutePath,
            operatorName = StatefulOperatorsUtils.STATE_STORE_SAVE_EXEC_OP_NAME
          )
        }
      }
    }
  }

  /**
   * Helper method to test SPARK-54420 read and write with different state format versions
   * for composite key aggregation (multiple grouping keys).
   * @param stateVersion The state format version (1 or 2)
   */
  private def testCompositeKeyRoundTripForStateVersion(stateVersion: Int): Unit = {
    withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> stateVersion.toString) {
      withTempDir { sourceDir =>
        withTempDir { targetDir =>
          // Step 1: Create state by running a composite key streaming aggregation
          runCompositeKeyStreamingAggregationQuery(sourceDir.getAbsolutePath)
          val inputData: MemoryStream[Int] = MemoryStream[Int]
          val aggregated = getCompositeKeyStreamingAggregationQuery(inputData)

          // add dummy data to target source to test writer won't load previous store
          testStream(aggregated, OutputMode.Update)(
            StartStream(checkpointLocation = targetDir.getAbsolutePath),
            // batch 0
            AddData(inputData, 0, 1),
            CheckLastBatch(
              (0, "Apple", 1, 0, 0, 0),
              (1, "Banana", 1, 1, 1, 1)
            )
          )

          // Perform round-trip test using common helper
          performRoundTripTest(
            sourceDir.getAbsolutePath,
            targetDir.getAbsolutePath,
            operatorName = StatefulOperatorsUtils.STATE_STORE_SAVE_EXEC_OP_NAME
          )
        }
      }
    }
  }

  /**
   * Helper method to test round-trip for stream-stream join with different versions.
   */
  private def testStreamStreamJoinRoundTrip(stateVersion: Int): Unit = {
    withSQLConf(SQLConf.STREAMING_JOIN_STATE_FORMAT_VERSION.key -> stateVersion.toString) {
      withTempDir { sourceDir =>
        withTempDir { targetDir =>
          // Step 1: Create state by running stream-stream join
          runStreamStreamJoinQuery(sourceDir.getAbsolutePath)

          // Create dummy data in target
          val inputData: MemoryStream[(Int, Long)] = MemoryStream[(Int, Long)]
          val query = getStreamStreamJoinQuery(inputData)
          testStream(query)(
            StartStream(checkpointLocation = targetDir.getAbsolutePath),
            AddData(inputData, (1, 1L)),
            CheckNewAnswer()
          )

          // Step 2: Test all 4 state stores created by stream-stream join
          val storeToColumnFamilies = StreamStreamJoinTestUtils.allStoreNames
            .map(s => s -> List(StateStore.DEFAULT_COL_FAMILY_NAME)).toMap
          // Perform round-trip test using common helper
          performRoundTripTest(
            sourceDir.getAbsolutePath,
            targetDir.getAbsolutePath,
            storeToColumnFamilies,
            operatorName = StatefulOperatorsUtils.SYMMETRIC_HASH_JOIN_EXEC_OP_NAME
          )
        }
      }
    }
  }

  /**
   * Helper method to test round-trip for flatMapGroupsWithState with different versions.
   */
  private def testFlatMapGroupsWithStateRoundTrip(stateVersion: Int): Unit = {
    // Skip this test on big endian platforms (version 1 only)
    if (stateVersion == 1) {
      assume(java.nio.ByteOrder.nativeOrder().equals(java.nio.ByteOrder.LITTLE_ENDIAN))
    }

    withSQLConf(SQLConf.FLATMAPGROUPSWITHSTATE_STATE_FORMAT_VERSION.key -> stateVersion.toString) {
      withTempDir { sourceDir =>
        withTempDir { targetDir =>
          // Step 1: Create state by running flatMapGroupsWithState
          runFlatMapGroupsWithStateQuery(sourceDir.getAbsolutePath)

          // Create dummy data in target
          val clock = new StreamManualClock
          val inputData: MemoryStream[(String, Long)] = MemoryStream[(String, Long)]
          val query = getFlatMapGroupsWithStateQuery(inputData)
          testStream(query, OutputMode.Update)(
            StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
              checkpointLocation = targetDir.getAbsolutePath),
            AddData(inputData, ("a", 1L)),
            AdvanceManualClock(1 * 1000),
            CheckLastBatch(("a", 1, 0, false))
          )

          // Perform round-trip test using common helper
          performRoundTripTest(
            sourceDir.getAbsolutePath,
            targetDir.getAbsolutePath,
            operatorName = StatefulOperatorsUtils.FLAT_MAP_GROUPS_WITH_STATE_EXEC_OP_NAME
          )
        }
      }
    }
  }

  /**
   * Helper method to build timer column family schemas and options for
   * RunningCountStatefulProcessorWithProcTimeTimer and EventTimeTimerProcessor
   *
   * @param timeMode Either TimeMode.EventTime() or TimeMode.ProcessingTime()
   * @return A tuple of three elements:
   *         - columnFamilyToSchemaMap: Maps column family names to their schema info
   *         - columnFamilyToSelectExprs: Maps column family names to custom select expressions
   *         - columnFamilyToStateSourceOptions: Maps column family names to state source options
   */
  private def getTimerStateConfigsForCountState(timeMode: TimeMode): (
      Map[String, StatePartitionWriterColumnFamilyInfo],
      Map[String, Seq[String]],
      Map[String, Map[String, String]]) = {

    val schemas = TimerTestUtils.getTimerConfigsForCountState(timeMode)
    val columnFamilyToSchemaMap = schemas.map { case (cfName, metadata) =>
      cfName -> createColFamilyInfo(
        metadata.keySchema,
        metadata.valueSchema,
        metadata.encoderSpec,
        cfName,
        metadata.useMultipleValuePerKey)
    }

    val (keyToTimestampCF, timestampToKeyCF) =
      TimerStateUtils.getTimerStateVarNames(timeMode.toString)

    val columnFamilyToSelectExprs = TimerTestUtils.getTimerColumnFamilyToSelectExprs(timeMode)

    val columnFamilyToStateSourceOptions = schemas.keys.map { cfName =>
        if (cfName == keyToTimestampCF || cfName == timestampToKeyCF) {
          cfName -> Map(StateSourceOptions.READ_REGISTERED_TIMERS -> "true")
        } else {
          cfName -> Map(StateSourceOptions.STATE_VAR_NAME -> cfName)
        }
    }.toMap

    (columnFamilyToSchemaMap, columnFamilyToSelectExprs, columnFamilyToStateSourceOptions)
  }


  // Run all tests with both changelog checkpointing enabled and disabled
  Seq(true, false).foreach { changelogCheckpointingEnabled =>
    val changelogCpTestSuffix = if (changelogCheckpointingEnabled) {
      "with changelog checkpointing"
    } else {
      "without changelog checkpointing"
    }

    def testWithChangelogConfig(testName: String)(testFun: => Unit): Unit = {
      test(s"$testName ($changelogCpTestSuffix)") {
        withSQLConf(
          "spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled" ->
            changelogCheckpointingEnabled.toString) {
          testFun
        }
      }
    }

    // Run transformWithState tests with enable/disable checkpoint V2
    Seq(1, 2).foreach { ckptVersion =>
      def testWithChangelogAndCheckpointId(testName: String)(testFun: => Unit): Unit = {
        testWithChangelogConfig(s"$testName (enableCkptId = ${ckptVersion >= 2})") {
          withSQLConf(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> ckptVersion.toString) {
            testFun
          }
        }
      }

      testWithChangelogAndCheckpointId("SPARK-54420: aggregation state ver 1") {
        testRoundTripForAggrStateVersion(1)
      }

      Seq(1, 2).foreach { version =>
        testWithChangelogAndCheckpointId(s"SPARK-54420: stream-stream join state ver $version") {
          testStreamStreamJoinRoundTrip(version)
        }
      }
      // Run transformWithState tests with different encoding formats
      Seq("unsaferow", "avro").foreach { encodingFormat =>
        def testWithChangelogAndEncodingConfig(testName: String)(testFun: => Unit): Unit = {
          testWithChangelogAndCheckpointId(
            s"$testName (encoding = $encodingFormat)") {
            withSQLConf(
              SQLConf.STREAMING_STATE_STORE_ENCODING_FORMAT.key -> encodingFormat) {
              testFun
            }
          }
        }

        testWithChangelogAndEncodingConfig(
          "SPARK-54411: transformWithState with multiple column families") {
          withTempDir { sourceDir =>
            withTempDir { targetDir =>
              val inputData = MemoryStream[String]
              val query = inputData.toDS()
                .groupByKey(x => x)
                .transformWithState(new MultiStateVarProcessor(),
                  TimeMode.None(),
                  OutputMode.Update())
              def runQuery(checkpointLocation: String, roundsOfData: Int): Unit = {
                val dataActions = (1 to roundsOfData).flatMap { _ =>
                  Seq(
                    AddData(inputData, "a", "b", "a"),
                    ProcessAllAvailable()
                  )
                }
                testStream(query)(
                  Seq(StartStream(checkpointLocation = checkpointLocation)) ++
                    dataActions ++
                    Seq(StopStream): _*
                )
              }

              runQuery(sourceDir.getAbsolutePath, 2)
              runQuery(targetDir.getAbsolutePath, 1)

              val schemas = MultiStateVarProcessorTestUtils.getSchemasWithMetadata()
              val columnFamilyToSelectExprs = MultiStateVarProcessorTestUtils
                .getColumnFamilyToSelectExprs()

              val columnFamilyToStateSourceOptions = schemas.keys.map { cfName =>
                val base = Map(
                  StateSourceOptions.STATE_VAR_NAME -> cfName
                )

                val withFlatten =
                  if (cfName == MultiStateVarProcessorTestUtils.ITEMS_LIST) {
                    base + (StateSourceOptions.FLATTEN_COLLECTION_TYPES -> "true")
                  } else {
                    base
                  }

                cfName -> withFlatten
              }.toMap

              performRoundTripTest(
                sourceDir.getAbsolutePath,
                targetDir.getAbsolutePath,
                storeToColumnFamilies = Map(StateStoreId.DEFAULT_STORE_NAME -> schemas.keys.toList),
                storeToColumnFamilyToSelectExprs =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> columnFamilyToSelectExprs),
                storeToColumnFamilyToStateSourceOptions =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> columnFamilyToStateSourceOptions),
                operatorName = StatefulOperatorsUtils.TRANSFORM_WITH_STATE_EXEC_OP_NAME
              )
            }
          }
        }

        testWithChangelogAndEncodingConfig(
          "SPARK-54411: transformWithState with eventTime timers") {
          withTempDir { sourceDir =>
            withTempDir { targetDir =>
              val inputData = MemoryStream[(String, Long)]
              val result = inputData.toDS()
                .select(col("_1").as("key"), timestamp_seconds(col("_2")).as("eventTime"))
                .withWatermark("eventTime", "10 seconds")
                .as[(String, Timestamp)]
                .groupByKey(_._1)
                .transformWithState(
                  new EventTimeTimerProcessor(),
                  TimeMode.EventTime(),
                  OutputMode.Update())

              testStream(result, OutputMode.Update())(
                StartStream(checkpointLocation = sourceDir.getAbsolutePath),
                AddData(inputData, ("a", 1L), ("b", 2L), ("c", 3L)),
                ProcessAllAvailable(),
                StopStream
              )

              testStream(result, OutputMode.Update())(
                StartStream(checkpointLocation = targetDir.getAbsolutePath),
                AddData(inputData, ("x", 1L)),
                ProcessAllAvailable(),
                StopStream
              )

              val (schemaMap, selectExprs, stateSourceOptions) =
                getTimerStateConfigsForCountState(TimeMode.EventTime())

              performRoundTripTest(
                sourceDir.getAbsolutePath,
                targetDir.getAbsolutePath,
                storeToColumnFamilies =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> schemaMap.keys.toList),
                storeToColumnFamilyToSelectExprs =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> selectExprs),
                storeToColumnFamilyToStateSourceOptions =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> stateSourceOptions),
                operatorName = StatefulOperatorsUtils.TRANSFORM_WITH_STATE_EXEC_OP_NAME
              )
            }
          }
        }

        testWithChangelogAndEncodingConfig(
          "SPARK-54411: transformWithState with processing time timers") {
          withTempDir { sourceDir =>
            withTempDir { targetDir =>
              val clock = new StreamManualClock
              val inputData = MemoryStream[String]
              val result = inputData.toDS()
                .groupByKey(x => x)
                .transformWithState(new RunningCountStatefulProcessorWithProcTimeTimer(),
                  TimeMode.ProcessingTime(),
                  OutputMode.Update())

              testStream(result, OutputMode.Update())(
                StartStream(checkpointLocation = sourceDir.getAbsolutePath,
                  trigger = Trigger.ProcessingTime("1 second"),
                  triggerClock = clock),
                AddData(inputData, "a"),
                AdvanceManualClock(1 * 1000),
                CheckNewAnswer(("a", "1")),
                StopStream
              )

              val clock2 = new StreamManualClock
              testStream(result, OutputMode.Update())(
                StartStream(checkpointLocation = targetDir.getAbsolutePath,
                  trigger = Trigger.ProcessingTime("1 second"),
                  triggerClock = clock2),
                AddData(inputData, "x"),
                AdvanceManualClock(1 * 1000),
                CheckNewAnswer(("a", "1"), ("x", "1")),
                StopStream
              )

              val (schemaMap, selectExprs, sourceOptions) =
                getTimerStateConfigsForCountState(TimeMode.ProcessingTime())

              performRoundTripTest(
                sourceDir.getAbsolutePath,
                targetDir.getAbsolutePath,
                storeToColumnFamilies =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> schemaMap.keys.toList),
                storeToColumnFamilyToSelectExprs =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> selectExprs),
                storeToColumnFamilyToStateSourceOptions =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> sourceOptions),
                operatorName = StatefulOperatorsUtils.TRANSFORM_WITH_STATE_EXEC_OP_NAME
              )
            }
          }
        }

        testWithChangelogAndEncodingConfig("SPARK-54411: transformWithState with list and TTL") {
          withTempDir { sourceDir =>
            withTempDir { targetDir =>
              val clock = new StreamManualClock
              val inputData = MemoryStream[InputEvent]
              val ttlConfig = TTLConfig(ttlDuration = Duration.ofMinutes(1))
              val result = inputData.toDS()
                .groupByKey(x => x.key)
                .transformWithState(new ListStateTTLProcessor(ttlConfig),
                  TimeMode.ProcessingTime(),
                  OutputMode.Update())

              testStream(result, OutputMode.Update())(
                StartStream(checkpointLocation = sourceDir.getAbsolutePath,
                  trigger = Trigger.ProcessingTime("1 second"),
                  triggerClock = clock),
                AddData(inputData, InputEvent("k1", "put", 1)),
                AdvanceManualClock(1 * 1000),
                CheckNewAnswer(),
                StopStream
              )

              val clock2 = new StreamManualClock
              testStream(result, OutputMode.Update())(
                StartStream(checkpointLocation = targetDir.getAbsolutePath,
                  trigger = Trigger.ProcessingTime("1 second"),
                  triggerClock = clock2),
                AddData(inputData, InputEvent("k1", "append", 1)),
                AdvanceManualClock(1 * 1000),
                CheckNewAnswer(),
                StopStream
              )

              val schemas = TTLProcessorUtils.getListStateTTLSchemasWithMetadata()

              val columnFamilyToSelectExprs = Map(
                TTLProcessorUtils.LIST_STATE -> TTLProcessorUtils.getTTLSelectExpressions(
                  TTLProcessorUtils.LIST_STATE
                ))

              val columnFamilyToStateSourceOptions = schemas.keys.map { cfName =>
                val base = Map(StateSourceOptions.STATE_VAR_NAME -> cfName)

                val withFlatten =
                  if (cfName == TTLProcessorUtils.LIST_STATE) {
                    base + (StateSourceOptions.FLATTEN_COLLECTION_TYPES -> "true")
                  } else {
                    base
                  }

                cfName -> withFlatten
              }.toMap

              performRoundTripTest(
                sourceDir.getAbsolutePath,
                targetDir.getAbsolutePath,
                storeToColumnFamilies =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> schemas.keys.toList),
                storeToColumnFamilyToSelectExprs =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> columnFamilyToSelectExprs),
                storeToColumnFamilyToStateSourceOptions =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> columnFamilyToStateSourceOptions),
                operatorName = StatefulOperatorsUtils.TRANSFORM_WITH_STATE_EXEC_OP_NAME
              )
            }
          }
        }

        testWithChangelogAndEncodingConfig("SPARK-54411: transformWithState with map and TTL") {
          withTempDir { sourceDir =>
            withTempDir { targetDir =>
              val clock = new StreamManualClock
              val inputData = MemoryStream[MapInputEvent]
              val ttlConfig = TTLConfig(ttlDuration = Duration.ofMinutes(1))
              val result = inputData.toDS()
                .groupByKey(x => x.key)
                .transformWithState(new MapStateTTLProcessor(ttlConfig),
                  TimeMode.ProcessingTime(),
                  OutputMode.Update())

              testStream(result)(
                StartStream(checkpointLocation = sourceDir.getAbsolutePath,
                  trigger = Trigger.ProcessingTime("1 second"),
                  triggerClock = clock),
                AddData(inputData, MapInputEvent("a", "key1", "put", 1)),
                AdvanceManualClock(1 * 1000),
                CheckNewAnswer(),
                StopStream
              )

              val clock2 = new StreamManualClock
              testStream(result)(
                StartStream(checkpointLocation = targetDir.getAbsolutePath,
                  trigger = Trigger.ProcessingTime("1 second"),
                  triggerClock = clock2),
                AddData(inputData, MapInputEvent("x", "key1", "put", 1)),
                AdvanceManualClock(1 * 1000),
                CheckNewAnswer(),
                StopStream
              )

              val schemas = TTLProcessorUtils.getMapStateTTLSchemasWithMetadata()

              val columnFamilyToSelectExprs = Map(
                TTLProcessorUtils.MAP_STATE -> TTLProcessorUtils.getTTLSelectExpressions(
                  TTLProcessorUtils.MAP_STATE
                ))

              val columnFamilyToStateSourceOptions = schemas.keys.map { cfName =>
                cfName -> Map(StateSourceOptions.STATE_VAR_NAME -> cfName)
              }.toMap

              performRoundTripTest(
                sourceDir.getAbsolutePath,
                targetDir.getAbsolutePath,
                storeToColumnFamilies =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> schemas.keys.toList),
                storeToColumnFamilyToSelectExprs =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> columnFamilyToSelectExprs),
                storeToColumnFamilyToStateSourceOptions =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> columnFamilyToStateSourceOptions),
                operatorName = StatefulOperatorsUtils.TRANSFORM_WITH_STATE_EXEC_OP_NAME
              )
            }
          }
        }

        testWithChangelogAndEncodingConfig("SPARK-54411: transformWithState with value and TTL") {
          withTempDir { sourceDir =>
            withTempDir { targetDir =>
              val clock = new StreamManualClock
              val inputData = MemoryStream[InputEvent]
              val ttlConfig = TTLConfig(ttlDuration = Duration.ofMinutes(1))
              val result = inputData.toDS()
                .groupByKey(x => x.key)
                .transformWithState(new ValueStateTTLProcessor(ttlConfig),
                  TimeMode.ProcessingTime(),
                  OutputMode.Update())

              testStream(result)(
                StartStream(checkpointLocation = sourceDir.getAbsolutePath,
                  trigger = Trigger.ProcessingTime("1 second"),
                  triggerClock = clock),
                AddData(inputData, InputEvent("k1", "put", 1)),
                AddData(inputData, InputEvent("k2", "put", 2)),
                AdvanceManualClock(1 * 1000),
                CheckNewAnswer(),
                StopStream
              )

              val clock2 = new StreamManualClock
              testStream(result)(
                StartStream(checkpointLocation = targetDir.getAbsolutePath,
                  trigger = Trigger.ProcessingTime("1 second"),
                  triggerClock = clock2),
                AddData(inputData, InputEvent("x", "put", 1)),
                AdvanceManualClock(1 * 1000),
                CheckNewAnswer(),
                StopStream
              )

              val schemas = TTLProcessorUtils.getValueStateTTLSchemasWithMetadata()

              val columnFamilyToStateSourceOptions = schemas.keys.map { cfName =>
                cfName -> Map(StateSourceOptions.STATE_VAR_NAME -> cfName)
              }.toMap

              performRoundTripTest(
                sourceDir.getAbsolutePath,
                targetDir.getAbsolutePath,
                storeToColumnFamilies =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> schemas.keys.toList),
                storeToColumnFamilyToStateSourceOptions =
                  Map(StateStoreId.DEFAULT_STORE_NAME -> columnFamilyToStateSourceOptions),
                operatorName = StatefulOperatorsUtils.TRANSFORM_WITH_STATE_EXEC_OP_NAME
              )
            }
          }
        }
      } // End of foreach loop for encoding format (transformWithState tests only)
    }

    testWithChangelogConfig("SPARK-54420: aggregation state ver 2") {
      testRoundTripForAggrStateVersion(2)
    }

    StreamingAggregationStateManager.supportedVersions.foreach { version =>
      testWithChangelogConfig(s"SPARK-54420: composite key aggregation state ver $version") {
        testCompositeKeyRoundTripForStateVersion(version)
      }
    }

    testWithChangelogConfig("SPARK-54420: dropDuplicatesWithinWatermark") {
      withTempDir { sourceDir =>
        withTempDir { targetDir =>
          // Step 1: Create state by running dropDuplicatesWithinWatermark
          runDropDuplicatesWithinWatermarkQuery(sourceDir.getAbsolutePath)

          // Create dummy data in target
          val inputData: MemoryStream[(String, Int)] = MemoryStream[(String, Int)]
          val deduped = getDropDuplicatesWithinWatermarkQuery(inputData)
          testStream(deduped, OutputMode.Append)(
            StartStream(checkpointLocation = targetDir.getAbsolutePath),
            AddData(inputData, ("a", 1)),
            CheckAnswer(("a", 1))
          )

          // Perform round-trip test using common helper
          performRoundTripTest(
            sourceDir.getAbsolutePath,
            targetDir.getAbsolutePath,
            operatorName = StatefulOperatorsUtils.DEDUPLICATE_WITHIN_WATERMARK_EXEC_OP_NAME
          )
        }
      }
    }

    testWithChangelogConfig("SPARK-54420: dropDuplicates with column specified") {
      withTempDir { sourceDir =>
        withTempDir { targetDir =>
          // Step 1: Create state by running dropDuplicates with column
          runDropDuplicatesQueryWithColumnSpecified(sourceDir.getAbsolutePath)

          // Create dummy data in target
          val inputData: MemoryStream[(String, Int)] = MemoryStream[(String, Int)]
          val deduped = getDropDuplicatesQueryWithColumnSpecified(inputData)
          testStream(deduped, OutputMode.Append)(
            StartStream(checkpointLocation = targetDir.getAbsolutePath),
            AddData(inputData, ("a", 1)),
            CheckAnswer(("a", 1))
          )

          // Perform round-trip test using common helper
          performRoundTripTest(
            sourceDir.getAbsolutePath,
            targetDir.getAbsolutePath,
            operatorName = StatefulOperatorsUtils.DEDUPLICATE_EXEC_OP_NAME
          )
        }
      }
    }

    testWithChangelogConfig("SPARK-54420: session window aggregation") {
      withTempDir { sourceDir =>
        withTempDir { targetDir =>
          // Step 1: Create state by running session window aggregation
          runSessionWindowAggregationQuery(sourceDir.getAbsolutePath)

          // Create dummy data in target
          val inputData: MemoryStream[(String, Long)] = MemoryStream[(String, Long)]
          val aggregated = getSessionWindowAggregationQuery(inputData)
          testStream(aggregated, OutputMode.Complete())(
            StartStream(checkpointLocation = targetDir.getAbsolutePath),
            AddData(inputData, ("a", 40L)),
            CheckNewAnswer(
              ("a", 40, 50, 10, 1)
            ),
            StopStream
          )

          // Perform round-trip test using common helper
          performRoundTripTest(
            sourceDir.getAbsolutePath,
            targetDir.getAbsolutePath,
            operatorName = StatefulOperatorsUtils.SESSION_WINDOW_STATE_STORE_SAVE_EXEC_OP_NAME
          )
        }
      }
    }

    testWithChangelogConfig("SPARK-54420: dropDuplicates") {
      withTempDir { sourceDir =>
        withTempDir { targetDir =>

          // Step 1: Create state by running a streaming aggregation
          runDropDuplicatesQuery(sourceDir.getAbsolutePath)
          val inputData: MemoryStream[Int] = MemoryStream[Int]
          val stream = getDropDuplicatesQuery(inputData)
          testStream(stream, OutputMode.Append)(
            StartStream(checkpointLocation = targetDir.getAbsolutePath),
            AddData(inputData, (1 to 5).flatMap(_ => (10 to 15)): _*),
            CheckAnswer(10 to 15: _*),
            assertNumStateRows(total = 6, updated = 6)
          )

          // Perform round-trip test using common helper
          performRoundTripTest(
            sourceDir.getAbsolutePath,
            targetDir.getAbsolutePath,
            operatorName = StatefulOperatorsUtils.DEDUPLICATE_EXEC_OP_NAME
          )
        }
      }
    }

    Seq(1, 2).foreach { version =>
      testWithChangelogConfig(s"SPARK-54420: flatMapGroupsWithState state ver $version") {
        testFlatMapGroupsWithStateRoundTrip(version)
      }
    }

    testWithChangelogConfig("SPARK-54411: stream-stream join state ver 3") {
      withSQLConf(
        SQLConf.STREAMING_JOIN_STATE_FORMAT_VERSION.key -> "3"
      ) {
        withTempDir { sourceDir =>
          withTempDir { targetDir =>
            val inputData = MemoryStream[(Int, Long)]
            val query = getStreamStreamJoinQuery(inputData)

            def runQuery(checkpointLocation: String, roundsOfData: Int): Unit = {
              val dataActions = (1 to roundsOfData).flatMap { _ =>
                Seq(
                  AddData(inputData, (1, 1L), (2, 2L), (3, 3L), (4, 4L), (5, 5L)),
                  ProcessAllAvailable()
                )
              }
              testStream(query)(
                (Seq(StartStream(checkpointLocation = checkpointLocation)) ++
                  dataActions ++
                  Seq(StopStream)): _*
              )
            }

            runQuery(sourceDir.getAbsolutePath, roundsOfData = 2)
            runQuery(targetDir.getAbsolutePath, roundsOfData = 1)

            val allColFamilyNames = StreamStreamJoinTestUtils.allStoreNames.toList
            performRoundTripTest(
              sourceDir.getAbsolutePath,
              targetDir.getAbsolutePath,
              storeToColumnFamilies = Map(StateStoreId.DEFAULT_STORE_NAME -> allColFamilyNames),
              storeToColumnFamilyToStateSourceOptions =
                Map(StateStoreId.DEFAULT_STORE_NAME -> allColFamilyNames.map {
                  cfName => cfName -> Map(StateSourceOptions.STORE_NAME -> cfName)
                }.toMap),
              operatorName = StatefulOperatorsUtils.SYMMETRIC_HASH_JOIN_EXEC_OP_NAME
            )
          }
        }
      }
    }
  } // End of foreach loop for changelog checkpointing dimension

  test("SPARK-54590: Rewriter throw exception if checkpoint version is not set correct") {
    withSQLConf(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "2") {
      withTempDir { sourceDir =>
        // Step 1: Create state by running a streaming aggregation
        runDropDuplicatesQuery(sourceDir.getAbsolutePath)
        val sourceCheckpointMetadata = new StreamingQueryCheckpointMetadata(
          spark, sourceDir.getAbsolutePath)
        val readBatchId = sourceCheckpointMetadata.commitLog.getLatestBatchId().get
        // Forced set STATE_STORE_CHECKPOINT_FORMAT_VERSION to 1 to mimic when user forgot to
        // update checkpoint version to 2 in sqlConfig when running stateRewriter
        // on checkpointV2 query.
        spark.conf.unset(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key)
        val ex = intercept[StateRewriterInvalidCheckpointError] {
          val rewriter = new StateRewriter(
            spark,
            readBatchId,
            readBatchId + 1,
            sourceDir.getAbsolutePath,
            spark.sessionState.newHadoopConf()
          )
          rewriter.run()
        }

        checkError(
          ex,
          "STATE_REWRITER_INVALID_CHECKPOINT.STATE_CHECKPOINT_FORMAT_VERSION_MISMATCH",
          "55019",
          parameters = Map(
            "checkpointLocation" -> sourceDir.getAbsolutePath,
            "expectedVersion" -> "2",
            "actualVersion" -> "1",
            "sqlConfKey" -> SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key)
        )
      }
    }
  }

}
