/*
 * 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, FileInputStream, InputStream}
import java.nio.charset.StandardCharsets.UTF_8
import java.nio.file.Files
import java.util.concurrent.ConcurrentHashMap
import java.util.zip.{ZipEntry, ZipOutputStream}

import scala.collection.{mutable, Map}
import scala.math._

import com.fasterxml.jackson.annotation.JsonInclude.Include
import com.fasterxml.jackson.databind.{DeserializationFeature, ObjectMapper}
import com.fasterxml.jackson.databind.annotation.JsonDeserialize
import com.fasterxml.jackson.module.scala.{ClassTagExtensions, DefaultScalaModule}
import org.apache.commons.io.FilenameUtils
import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.fs.{FileStatus, FileSystem, Path, PathFilter}
import org.json4s.{Formats, NoTypeHints}
import org.json4s.jackson.Serialization

import org.apache.spark.{SparkConf, SparkEnv, SparkException}
import org.apache.spark.internal.{Logging, LogKeys, MessageWithContext}
import org.apache.spark.internal.LogKeys.{DFS_FILE, VERSION_NUM}
import org.apache.spark.io.CompressionCodec
import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.execution.streaming.checkpointing.{CheckpointFileManager, ChecksumCheckpointFileManager, ChecksumFile}
import org.apache.spark.util.ArrayImplicits._
import org.apache.spark.util.Utils

/**
 * Class responsible for syncing RocksDB checkpoint files from local disk to DFS.
 * For each version, checkpoint is saved in specific directory structure that allows successive
 * versions to reuse to SST data files and archived log files. This allows each commit to be
 * incremental, only new SST files and archived log files generated by RocksDB will be uploaded.
 * The directory structures on local disk and in DFS are as follows.
 *
 * Local checkpoint dir structure
 * ------------------------------
 * RocksDB generates a bunch of files in the local checkpoint directory. The most important among
 * them are the SST files; they are the actual log structured data files. Rest of the files contain
 * the metadata necessary for RocksDB to read the SST files and start from the checkpoint.
 * Note that the SST files are hard links to files in the RocksDB's working directory, and therefore
 * successive checkpoints can share some of the SST files. So these SST files have to be copied to
 * DFS in shared directory such that different committed versions can save them.
 *
 * We consider both SST files and archived log files as immutable files which can be shared between
 * different checkpoints.
 *
 *    localCheckpointDir
 *                  |
 *                  +-- OPTIONS-000005
 *                  +-- MANIFEST-000008
 *                  +-- CURRENT
 *                  +-- 00007.sst
 *                  +-- 00011.sst
 *                  +-- archive
 *                  |     +-- 00008.log
 *                  |     +-- 00013.log
 *                  ...
 *
 *
 * DFS directory structure after saving to DFS as version 10
 * -----------------------------------------------------------
 * The SST and archived log files are given unique file names and copied to the shared subdirectory.
 * Every version maintains a mapping of local immutable file name to the unique file name in DFS.
 * This mapping is saved in a JSON file (named `metadata`), which is zipped along with other
 * checkpoint files into a single file `[version].zip`.
 *
 *    dfsRootDir
 *           |
 *           +-- SSTs
 *           |      +-- 00007-[uuid1].sst
 *           |      +-- 00011-[uuid2].sst
 *           +-- logs
 *           |      +-- 00008-[uuid3].log
 *           |      +-- 00013-[uuid4].log
 *           +-- 10.zip
 *           |      +-- metadata         <--- contains mapping between 00007.sst and [uuid1].sst,
 *                                            and the mapping between 00008.log and [uuid3].log
 *           |      +-- OPTIONS-000005
 *           |      +-- MANIFEST-000008
 *           |      +-- CURRENT
 *           |      ...
 *           |
 *           +-- 9.zip
 *           +-- 8.zip
 *           ...
 *
 * Note the following.
 * - Each [version].zip is a complete description of all the data and metadata needed to recover
 *   a RocksDB instance at the corresponding version. The SST files and log files are not included
 *   in the zip files, they can be shared cross different versions. This is unlike the
 *   [version].delta files of HDFSBackedStateStore where previous delta files needs to be read
 *   to be recovered.
 * - This is safe wrt speculatively executed tasks running concurrently in different executors
 *   as each task would upload a different copy of the generated immutable files and
 *   atomically update the [version].zip.
 * - Immutable files are identified uniquely based on their file name and file size.
 * - Immutable files can be reused only across adjacent checkpoints/versions.
 * - This class is thread-safe. Specifically, it is safe to concurrently delete old files from a
 *   different thread than the task thread saving files.
 *
 * @param dfsRootDir  Directory where the [version].zip files will be stored
 * @param localTempDir Local directory for temporary work
 * @param hadoopConf   Hadoop configuration for talking to DFS
 * @param loggingId    Id that will be prepended in logs for isolating concurrent RocksDBs
 * @param fileChecksumEnabled Whether file checksum generation and verification is enabled
 * @param fileChecksumThreadPoolSize Number of threads used to concurrently operate on the
 *                                   main and checksum files
 */
class RocksDBFileManager(
    dfsRootDir: String,
    localTempDir: File,
    hadoopConf: Configuration,
    codecName: String = CompressionCodec.ZSTD,
    loggingId: String = "",
    storeConf: StateStoreConf = StateStoreConf.empty,
    fileChecksumEnabled: Boolean = false,
    fileChecksumThreadPoolSize: Option[Int] = None)
  extends Logging {

  import RocksDBImmutableFile._

  protected def getFileSystem(myDfsRootDir: String, myHadoopConf: Configuration) : FileSystem = {
    new Path(myDfsRootDir).getFileSystem(myHadoopConf)
  }

  private lazy val fm = {
    val mgr = CheckpointFileManager.create(new Path(dfsRootDir), hadoopConf)
    if (fileChecksumEnabled) {
      new ChecksumCheckpointFileManager(
        mgr,
        // Allowing this for perf, since we do orphan checksum file cleanup in maintenance anyway
        allowConcurrentDelete = true,
        numThreads = fileChecksumThreadPoolSize.get,
        skipCreationIfFileMissingChecksum
          = storeConf.checkpointFileChecksumSkipCreationIfFileMissingChecksum)
    } else {
      mgr
    }
  }

  private val fs = getFileSystem(dfsRootDir, hadoopConf)
  private val onlyZipFiles = new PathFilter {
    override def accept(path: Path): Boolean = path.toString.endsWith(".zip")
  }
  private val onlyChangelogFiles = new PathFilter {
    override def accept(path: Path): Boolean = path.toString.endsWith(".changelog")
  }

  private lazy val sparkConf = Option(SparkEnv.get).map(_.conf).getOrElse(new SparkConf)

  private def codec = CompressionCodec.createCodec(sparkConf, codecName)

  // This is set when a version is loaded/committed. Hence only set by a task thread.
  private var maxSeenVersion: Option[Long] = None
  // This is set during deletion of old versions. Hence only set by a maintenance thread.
  private var minSeenVersion = 1L

  @volatile private var rootDirChecked: Boolean = false

  // (version, checkpointUniqueId) -> immutable files
  private val versionToRocksDBFiles =
    new ConcurrentHashMap[(Long, Option[String]), Seq[RocksDBImmutableFile]]()

  /** Close this file manager and release underlying resources. */
  def close(): Unit = {
    fm.close()
  }

  /**
   * Get the changelog version based on rocksDB features.
   * @return the version of changelog
   */
  private def getChangelogWriterVersion(
      useColumnFamilies: Boolean,
      stateStoreCheckpointIdEnabled: Boolean): Short = {
    (useColumnFamilies, stateStoreCheckpointIdEnabled) match {
      case (false, false) => 1
      case (true, false) => 2
      case (false, true) => 3
      case _ => 4
    }
  }

  private def createDfsRootDirIfNotExist(): Unit = {
    if (!rootDirChecked) {
      val rootDir = new Path(dfsRootDir)
      if (!fm.exists(rootDir)) fm.mkdirs(rootDir)
      rootDirChecked = true
    }
  }

  def getChangeLogWriter(
      version: Long,
      useColumnFamilies: Boolean = false,
      checkpointUniqueId: Option[String] = None,
      stateStoreCheckpointIdLineage: Option[Array[LineageItem]] = None
    ): StateStoreChangelogWriter = {
    try {
      val changelogFile = dfsChangelogFile(version, checkpointUniqueId)
      createDfsRootDirIfNotExist()

      val enableStateStoreCheckpointIds = checkpointUniqueId.isDefined
      val changelogVersion = getChangelogWriterVersion(
        useColumnFamilies, enableStateStoreCheckpointIds)

      val changelogWriter = changelogVersion match {
        case 1 =>
          new StateStoreChangelogWriterV1(fm, changelogFile, codec)
        case 2 =>
          new StateStoreChangelogWriterV2(fm, changelogFile, codec)
        case 3 =>
          assert(enableStateStoreCheckpointIds && stateStoreCheckpointIdLineage.isDefined,
            "StateStoreChangelogWriterV3 should only be initialized when " +
              "state store checkpoint unique id is enabled")
          new StateStoreChangelogWriterV3(fm, changelogFile, codec,
            stateStoreCheckpointIdLineage.get)
        case 4 =>
          assert(enableStateStoreCheckpointIds && stateStoreCheckpointIdLineage.isDefined,
            "StateStoreChangelogWriterV4 should only be initialized when " +
              "state store checkpoint unique id is enabled")
          new StateStoreChangelogWriterV4(fm, changelogFile, codec,
            stateStoreCheckpointIdLineage.get)
        case _ =>
          throw QueryExecutionErrors.invalidChangeLogWriterVersion(changelogVersion)
      }

      logInfo(log"Loaded change log reader version " +
        log"${MDC(LogKeys.FILE_VERSION, changelogWriter.version)}")

      changelogWriter
    } catch {
      case e: SparkException
        if Option(e.getCondition).exists(_.contains("CANNOT_LOAD_STATE_STORE")) =>
          throw e
      case e: Throwable =>
        throw StateStoreErrors.failedToGetChangelogWriter(version, e)
    }
  }

  // Get the changelog file at version
  def getChangelogReader(
      version: Long,
      checkpointUniqueId: Option[String] = None): StateStoreChangelogReader = {
    val changelogFile = dfsChangelogFile(version, checkpointUniqueId)
    val reader = new StateStoreChangelogReaderFactory(fm, changelogFile, codec)
      .constructChangelogReader()

    logInfo(log"Loaded change log reader version ${MDC(LogKeys.FILE_VERSION, reader.version)}")

    reader
  }

  /**
   * Metrics for loading checkpoint from DFS. Every loadCheckpointFromDFS call will update this
   * metrics, so this effectively records the latest metrics.
   */
  @volatile private var loadCheckpointMetrics = RocksDBFileManagerMetrics.EMPTY_METRICS

  /**
   * Metrics for saving checkpoint to DFS. Every saveCheckpointToDFS call will update this
   * metrics, so this effectively records the latest metrics.
   */
  @volatile private var saveCheckpointMetrics = RocksDBFileManagerMetrics.EMPTY_METRICS

  def latestLoadCheckpointMetrics: RocksDBFileManagerMetrics = loadCheckpointMetrics

  def latestSaveCheckpointMetrics: RocksDBFileManagerMetrics = saveCheckpointMetrics

  /** Save all the files in given local checkpoint directory as a committed version in DFS */
  def saveCheckpointToDfs(
      checkpointDir: File,
      version: Long,
      numKeys: Long,
      numInternalKeys: Long,
      fileMapping: Map[String, RocksDBSnapshotFile],
      columnFamilyMapping: Option[Map[String, ColumnFamilyInfo]] = None,
      maxColumnFamilyId: Option[Short] = None,
      checkpointUniqueId: Option[String] = None,
      verifyNonEmptyFilesInZip: Boolean = false): Unit = {
    logFilesInDir(checkpointDir, log"Saving checkpoint files " +
      log"for version ${MDC(LogKeys.VERSION_NUM, version)}")
    val (localImmutableFiles, localOtherFiles) = listRocksDBFiles(checkpointDir)
    val rocksDBFiles = saveImmutableFilesToDfs(
      version, localImmutableFiles, fileMapping, checkpointUniqueId)

    val colFamilyIdMapping: Option[Map[String, Short]] = if (columnFamilyMapping.isDefined) {
      Some(columnFamilyMapping.get.map {
        case (cfName, cfInfo) =>
          cfName -> cfInfo.cfId
      })
    } else {
      None
    }

    val colFamilyTypeMapping: Option[Map[String, Boolean]] = if (columnFamilyMapping.isDefined) {
      Some(columnFamilyMapping.get.map {
        case (cfName, cfInfo) =>
          cfName -> cfInfo.isInternal
      })
    } else {
      None
    }

    val metadata = RocksDBCheckpointMetadata(rocksDBFiles, numKeys, numInternalKeys,
      colFamilyIdMapping, colFamilyTypeMapping, maxColumnFamilyId)
    val metadataFile = localMetadataFile(checkpointDir)
    metadata.writeToFile(metadataFile)
    logDebug(log"Written metadata for version ${MDC(LogKeys.VERSION_NUM, version)}:\n" +
      log"${MDC(LogKeys.METADATA_JSON, metadata.prettyJson)}")

    val (_, zipFilesTimeMs) = Utils.timeTakenMs {
      if (version <= 1 && numKeys <= 0) {
        // If we're writing the initial version and there's no data, we have to initialize
        // the root directory. Normally saveImmutableFilesToDfs will do this initialization, but
        // when there's no data that method won't write any files, and zipToDfsFile uses the
        // CheckpointFileManager.createAtomic API which doesn't auto-initialize parent directories.
        // Moreover, once we disable to track the number of keys, in which the numKeys is -1, we
        // still need to create the initial dfs root directory anyway.
        createDfsRootDirIfNotExist()
      }
      zipToDfsFile(localOtherFiles :+ metadataFile,
        dfsBatchZipFile(version, checkpointUniqueId), verifyNonEmptyFilesInZip)
      logInfo(log"Saved checkpoint file for version ${MDC(LogKeys.VERSION_NUM, version)} " +
        log"checkpointUniqueId: ${MDC(LogKeys.UUID, checkpointUniqueId.getOrElse(""))}")
    }

    // populate the SaveCheckpointMetrics
    saveCheckpointMetrics =
      saveCheckpointMetrics.copy(
        // Round up to 1ms to reassure that we've logged successfully and avoid flaky tests
        saveZipFilesTimeMs = Some(Math.max(zipFilesTimeMs, 1L))
      )
  }

  /**
   * Load all necessary files for specific checkpoint version from DFS to given local directory.
   * If version is 0, then it will delete all files in the directory. For other versions, it
   * ensures that only the exact files generated during checkpointing will be present in the
   * local directory.
   */
  def loadCheckpointFromDfs(
      version: Long,
      localDir: File,
      rocksDBFileMapping: RocksDBFileMapping,
      checkpointUniqueId: Option[String] = None): RocksDBCheckpointMetadata = {
    logInfo(log"Loading checkpoint files for version ${MDC(LogKeys.VERSION_NUM, version)} " +
      log"checkpointUniqueId: ${MDC(LogKeys.UUID, checkpointUniqueId.getOrElse(""))}")
    // The unique ids of SST files are checked when opening a rocksdb instance. The SST files
    // in larger versions can't be reused even if they have the same size and name because
    // they belong to another rocksdb instance.
    versionToRocksDBFiles.keySet().removeIf(_._1 >= version)
    val metadata = if (version == 0) {
      if (localDir.exists) Utils.deleteRecursively(localDir)
      Utils.createDirectory(localDir)
      createDfsRootDirIfNotExist()
      // Since we cleared the local dir, we should also clear the local file mapping
      rocksDBFileMapping.clear()
      // Set empty metrics since we're not loading any files from DFS
      loadCheckpointMetrics = RocksDBFileManagerMetrics.EMPTY_METRICS
      RocksDBCheckpointMetadata(Seq.empty, 0)
    } else {
      // Delete all non-immutable files in local dir, and unzip new ones from DFS commit file
      listRocksDBFiles(localDir)._2.foreach(_.delete())
      unzipBatchZipFileFromDfs(version, checkpointUniqueId, localDir)

      // Copy the necessary immutable files
      val metadataFile = localMetadataFile(localDir)
      val metadata = RocksDBCheckpointMetadata.readFromFile(metadataFile)
      logDebug(log"Read metadata for version ${MDC(LogKeys.VERSION_NUM, version)}:\n" +
        log"${MDC(LogKeys.METADATA_JSON, metadata.prettyJson)}")
      loadImmutableFilesFromDfs(metadata.immutableFiles, localDir, rocksDBFileMapping, version)
      versionToRocksDBFiles.put((version, checkpointUniqueId), metadata.immutableFiles)
      metadataFile.delete()
      metadata
    }
    logFilesInDir(localDir, log"Loaded checkpoint files " +
      log"for version ${MDC(LogKeys.VERSION_NUM, version)}")
    logInfo(log"RocksDB file mapping after loading checkpoint version " +
      log"${MDC(LogKeys.VERSION_NUM, version)} from DFS:\n" +
      log"${MDC(LogKeys.ROCKS_DB_FILE_MAPPING, rocksDBFileMapping)}")

    metadata
  }

  // Return if there is a snapshot file at the corresponding version
  // and optionally with checkpointunique id, e.g. version.zip or version_uniqueId.zip
  def existsSnapshotFile(version: Long, checkpointUniqueId: Option[String] = None): Boolean = {
    createDfsRootDirIfNotExist()
    fm.exists(dfsBatchZipFile(version, checkpointUniqueId))
  }

  // Get latest snapshot version <= version
  def getLatestSnapshotVersion(version: Long): Long = {
    val path = new Path(dfsRootDir)
    if (fm.exists(path)) {
      // If the latest version snapshot exists, we avoid listing.
      if (fm.exists(dfsBatchZipFile(version))) {
        return version
      }
      fm.list(path, onlyZipFiles)
        .map(_.getPath.getName.stripSuffix(".zip"))
        .map(_.toLong)
        .filter(_ <= version)
        .foldLeft(0L)(math.max)
    } else {
      0
    }
  }

  /** Get the latest snapshot version available in DFS. If none present, it returns 0. */
  def getLatestSnapshotVersion(): Long = {
    val path = new Path(dfsRootDir)
    if (fm.exists(path)) {
      val files = fm.list(path).map(_.getPath)
      files.filter(onlyZipFiles.accept)
        .map { fileName =>
          fileName.getName.stripSuffix(".zip").split("_") match {
            case Array(version, _) => version.toLong
            case Array(version) => version.toLong
          }
        }.foldLeft(0L)(math.max)
    } else {
      0
    }
  }

  /** Get all the snapshot versions that can be used to load this version */
  def getEligibleSnapshotsForVersion(version: Long): Seq[Long] = {
    SnapshotLoaderHelper.getEligibleSnapshotsForVersion(
      version, fm, new Path(dfsRootDir), onlyZipFiles, fileSuffix = ".zip")
  }

  /**
   * Based on the ground truth lineage loaded from changelog file (lineage), this function
   * does file listing to find all snapshot (version, uniqueId) pairs, and finds
   * all snapshots that match the lineage, sorted by version descending.
   *
   * @param lineage The ground truth lineage loaded from changelog file, sorted by id
   * @return All snapshot (version, uniqueId) pairs matching the lineage, sorted descending
   *         by version. Empty if no matching snapshots found.
   */
  def getSnapshotVersionsAndUniqueIdsFromLineage(
      lineage: Array[LineageItem]): Seq[(Long, String)] = {
    val path = new Path(dfsRootDir)
    if (fm.exists(path)) {
      fm.list(path, onlyZipFiles)
        .map(_.getPath.getName.stripSuffix(".zip").split("_"))
        .collect {
          case Array(ver, id) if lineage.contains(LineageItem(ver.toLong, id)) =>
            (ver.toLong, id)
        }
        .sortBy(_._1)
        .reverse
        .toSeq
    } else {
      Seq.empty
    }
  }

  /**
   * Based on the ground truth lineage loaded from changelog file (lineage), this function
   * does file listing to find all snapshot (version, uniqueId) pairs, and finds
   * the ground truth latest snapshot (version, uniqueId) the db instance needs to load.
   *
   * @param lineage The ground truth lineage loaded from changelog file, sorted by id
   * @return The ground truth latest snapshot (version, uniqueId) the db instance needs to load,
   *         when the return value is None it means ther is no such snapshot found.
   */
  def getLatestSnapshotVersionAndUniqueIdFromLineage(
      lineage: Array[LineageItem]): Option[(Long, String)] = {
    getSnapshotVersionsAndUniqueIdsFromLineage(lineage).headOption
  }

  /** Get the latest version available in the DFS directory. If no data present, it returns 0. */
  def getLatestVersion(): Long = {
    val path = new Path(dfsRootDir)
    if (fm.exists(path)) {
      val files = fm.list(path).map(_.getPath)
      val changelogFileVersions = files.filter(onlyChangelogFiles.accept)
        .map { fileName =>
          fileName.getName.stripSuffix(".changelog").split("_") match {
            case Array(version, _) => version.toLong
            case Array(version) => version.toLong
          }
        }
      val snapshotFileVersions = files.filter(onlyZipFiles.accept)
        .map { fileName =>
          fileName.getName.stripSuffix(".zip").split("_") match {
            case Array(version, _) => version.toLong
            case Array(version) => version.toLong
          }
        }
      val versions = changelogFileVersions ++ snapshotFileVersions
      versions.foldLeft(0L)(math.max)
    } else {
      0
    }
  }

  /**
   * Find orphan files which are not tracked by zip files.
   * Both sst files and log files can be orphan files.
   * They are uploaded separately before the zip file of that version is uploaded.
   * When the zip file of a version get overwritten, the referenced sst and log files become orphan.
   * Be careful here since sst and log files of the ongoing version
   * also appear to be orphan before their zip file is uploaded.
   *
   * @param trackedFiles files tracked by metadata in versioned zip file
   * @param allFiles all sst or log files in the directory.
   * @return filenames of orphan files
   */
  def findOrphanFiles(trackedFiles: Seq[String], allFiles: Seq[FileStatus]): Seq[String] = {
    val fileModificationTimes = allFiles.map(file =>
      file.getPath.getName -> file.getModificationTime).toMap
    if (trackedFiles.nonEmpty && allFiles.size > trackedFiles.size) {
      // Some tracked files may not be in the directory when listing.
      val oldestTrackedFileModificationTime = trackedFiles.flatMap(fileModificationTimes.get(_)).min
      // If this immutable file is older than any tracked file,
      // then it can't belong to the ongoing version and it should be safe to clean it up.
      val orphanFiles = fileModificationTimes
        .filter(_._2 < oldestTrackedFileModificationTime).keys.toSeq
      if (orphanFiles.nonEmpty) {
        logInfo(log"Found ${MDC(LogKeys.NUM_FILES, orphanFiles.size)} orphan files: " +
          log"${MDC(LogKeys.FILE_MODIFICATION_TIME, orphanFiles.take(20).mkString(", "))}" +
          log"... (display at most 20 filenames) that should be deleted.")
      }
      orphanFiles
    } else {
      Seq.empty
    }
  }

  private def deleteChangelogFiles(
      versionsAndUniqueIdsToDelete: Array[(Long, Option[String])]): Unit = {
    versionsAndUniqueIdsToDelete.foreach { case (version, uniqueId) =>
      try {
        fm.delete(dfsChangelogFile(version, uniqueId))
        logInfo(log"Deleted changelog file ${MDC(LogKeys.VERSION_NUM, version)} uniqueId: " +
          log"${MDC(LogKeys.UUID, uniqueId.getOrElse(""))}")
      } catch {
        case e: Exception =>
          logWarning(
            log"Error deleting changelog file for version ${MDC(LogKeys.FILE_VERSION, version)} " +
              log"uniqueId: ${MDC(LogKeys.UUID, uniqueId.getOrElse(""))}", e)
      }
    }
  }

  /**
   *  Set maxSeenVersion to max of itself and version we are uploading.
   *  This is to ensure accuracy in the case the query has restarted from a particular version.
   */
  def setMaxSeenVersion(version: Long): Unit = {
    if (maxSeenVersion.isDefined) {
      maxSeenVersion = Some(Math.max(maxSeenVersion.get, version))
    } else {
      maxSeenVersion = Some(version)
    }
  }

  /**
   * Determines whether batch deletion of stale version files should be skipped
   * based on the following parameters and estimates of maximum and minimum
   * versions present in the checkpoint directory.
   *
   * @param numVersionsToRetain Number of versions to retain for rollbacks.
   * @param minVersionsToDelete Minimum number of stale versions required to trigger deletion.
   * @return `true` if insufficient stale versions present, otherwise `false`.
   */
  private def shouldSkipDeletion(numVersionsToRetain: Int, minVersionsToDelete: Long): Boolean = {
    // If minVersionsToDelete <= 0, we call list every time maintenance is invoked
    // This is the original behaviour without list api call optimization
    if (minVersionsToDelete > 0) {
      // When maxSeenVersion is defined, we check the if number of stale version files
      // are at least the value of minVersionsToDelete for batch deletion of files
      // We still proceed with deletion if maxSeenVersion isn't set to ensure the fallback
      // is to clean up files if maxSeenVersion fails to be initialized
      if (maxSeenVersion.isDefined) {
        logInfo(log"Estimated maximum version is " +
          log"${MDC(LogKeys.MAX_SEEN_VERSION, maxSeenVersion.get)}" +
          log" and minimum version is ${MDC(LogKeys.MIN_SEEN_VERSION, minSeenVersion)}")
        // If the number of versions to delete is negative, that means that none of the versions
        // are eligible for deletion and we set the variable to 0
        val versionsToDelete = max(maxSeenVersion.get - minSeenVersion + 1 - numVersionsToRetain, 0)
        if (versionsToDelete < minVersionsToDelete) {
          logInfo(log"Skipping deleting files." +
            log" Need at least ${MDC(LogKeys.MIN_VERSIONS_TO_DELETE, minVersionsToDelete)}" +
            log" stale versions for batch deletion but found only" +
            log" ${MDC(LogKeys.VERSIONS_TO_DELETE, versionsToDelete)}.")
          return true
        }
      }
    }
    false
  }

  /**
   * Delete old versions by deleting the associated version and SST files.
   * At a high-level, when enough stale version files are present for batch deletion,
   * this method finds which versions to delete, and which SST files that were
   * last used in those versions. It's safe to delete these SST files because a SST file can
   * be reused only in successive versions. Therefore, if a SST file F was last used in version
   * V, then it won't be used in version V+1 or later, and if version V can be deleted, then
   * F can safely be deleted as well.
   *
   * First, it checks whether enough stale version files are present for batch deletion.
   * If true, it does the following to find old files.
   * - List all the existing [version].zip files
   * - Find the min version that needs to be retained based on the given `numVersionsToRetain`.
   * - Accordingly decide which versions should be deleted.
   * - Resolve all SSTs files of all the existing versions, if not already resolved.
   * - Find the files that were last used in the to-be-deleted versions as we will not
   *   need those files any more.
   * - Find the orphan sst and log files whose zip files are not uploaded successfully
   *   or have been overwritten. To avoid deleting files of ongoing tasks, only delete orphan files
   *   that are older than all tracked files when there are at least 2 versions.
   * - Delete sst and log files in to-be-deleted versions.
   * - Delete orphan files.
   * - Delete changelog files of to-be-deleted versions.
   *
   * Note that it only deletes files that it knows are safe to delete.
   * It may not delete the following files.
   * - Partially written SST files
   * - SST files that were used in a version, but that version got overwritten with a different
   *   set of SST files.
   *
   * @param numVersionsToRetain the number of RocksDB versions to keep in object store after the
   *                            deletion. Must be greater than 0, or -1 to retain all versions.
   * @param maxVersionsToDeletePerMaintenance the max number of RocksDB versions
   *                            to delete per maintenance operation.
   *                            Must be greater than 0, or -1 to delete all stale versions.
   * @param minVersionsToDelete the min number of stale versions required to trigger deletion.
   *                            If its set to <= 0, then we will always perform list operations
   *                            to determine deletion candidates. If set to a positive value, then
   *                            we will skip deletion if the number of stale versions is less than
   *                            this value.
   */
  def deleteOldVersions(
      numVersionsToRetain: Int,
      maxVersionsToDeletePerMaintenance: Int = -1,
      minVersionsToDelete: Long = 0): Unit = {
    // Check if enough stale version files present
    if (shouldSkipDeletion(numVersionsToRetain, minVersionsToDelete)) return

    val path = new Path(dfsRootDir)
    val allFiles = fm.list(path).map(_.getPath)
    val snapshotFiles = allFiles.filter(file => onlyZipFiles.accept(file))
    val changelogFiles = allFiles.filter(file => onlyChangelogFiles.accept(file))
    // All versions present in DFS, sorted
    val sortedSnapshotVersionsAndUniqueIds = snapshotFiles
      .map(_.getName.stripSuffix(".zip").split("_"))
      .map {
        case Array(version, uniqueId) => (version.toLong, Some(uniqueId))
        case Array(version) => (version.toLong, None)
      }
      .sortBy(_._1)

    // Return if no versions generated yet
    if (sortedSnapshotVersionsAndUniqueIds.isEmpty) return

    // Find the versions to delete
    val maxSnapshotVersionPresent = sortedSnapshotVersionsAndUniqueIds.last._1
    val minSnapshotVersionPresent = sortedSnapshotVersionsAndUniqueIds.head._1

    // In order to reconstruct numVersionsToRetain version, retain the latest snapshot
    // that satisfies (version <= maxSnapshotVersionPresent - numVersionsToRetain + 1).
    // Also require
    // minVersionToRetain <= minSnapshotVersionPresent + maxVersionsToDeletePerMaintenance.
    // If none of the snapshots satisfy the condition, minVersionToRetain will be 0 and
    // no version gets deleted.
    val minVersionToRetain = sortedSnapshotVersionsAndUniqueIds
      .map(_._1)
      .filter(_ <= maxSnapshotVersionPresent - numVersionsToRetain + 1)
      .filter( v =>
        if (maxVersionsToDeletePerMaintenance != -1) {
          v <= minSnapshotVersionPresent + maxVersionsToDeletePerMaintenance
        } else true
      )
      .foldLeft(0L)(math.max)

    // When snapshotVersionToDelete is non-empty, there are at least 2 snapshot versions.
    // We only delete orphan files when there are at least 2 versions,
    // which avoid deleting files for running tasks.
    val snapshotVersionsAndUniqueIdsToDelete = sortedSnapshotVersionsAndUniqueIds
      .filter(_._1 < minVersionToRetain)
    val snapshotVersionsToDelete = snapshotVersionsAndUniqueIdsToDelete.map(_._1)
    if (snapshotVersionsAndUniqueIdsToDelete.isEmpty) return

    // Resolve RocksDB files for all the versions and find the max version each file is used
    val fileToMaxUsedVersion = new mutable.HashMap[String, Long]
    sortedSnapshotVersionsAndUniqueIds.foreach { case (version, uniqueId) =>
      var readFromCache = true
      val cachedFiles = Option(versionToRocksDBFiles.get((version, uniqueId))).getOrElse {
        readFromCache = false
        val newResolvedFiles = getImmutableFilesFromVersionZip(version, uniqueId)
        versionToRocksDBFiles.put((version, uniqueId), newResolvedFiles)
        newResolvedFiles
      }
      cachedFiles.foreach(f => fileToMaxUsedVersion(f.dfsFileName) =
        math.max(version, fileToMaxUsedVersion.getOrElse(f.dfsFileName, version)))

      // For ckpt v1, for the min retained version, fetch metadata from cloud zip
      // to protect files that may differ from the in-memory cache. This handles the case
      // where a no-overwrite FS prevented a snapshot zip from being overwritten: the
      // in-memory cache may reference a newer set of SST files (from a retry), while the
      // cloud zip still references the original SST files.
      // We do this for the min retained version, to make sure the files from the original
      // upload for this version are not seen as orphan files.
      if (uniqueId.isEmpty && version == minVersionToRetain && readFromCache) {
        val cloudFiles = getImmutableFilesFromVersionZip(version, uniqueId)
        cloudFiles.foreach(f => fileToMaxUsedVersion(f.dfsFileName) =
          math.max(version, fileToMaxUsedVersion.getOrElse(f.dfsFileName, version)))
      }
    }

    // Best effort attempt to delete SST files that were last used in to-be-deleted versions
    val filesToDelete = fileToMaxUsedVersion.filter {
      case (_, v) => snapshotVersionsToDelete.contains(v)
    }

    val sstDir = new Path(dfsRootDir, RocksDBImmutableFile.SST_FILES_DFS_SUBDIR)
    val logDir = new Path(dfsRootDir, RocksDBImmutableFile.LOG_FILES_DFS_SUBDIR)
    val allSstFiles = if (fm.exists(sstDir)) fm.list(sstDir).toImmutableArraySeq else Seq.empty
    val allLogFiles = if (fm.exists(logDir)) fm.list(logDir).toImmutableArraySeq else Seq.empty
    filesToDelete ++= findOrphanFiles(fileToMaxUsedVersion.keys.toSeq, allSstFiles ++ allLogFiles)
      .map(_ -> -1L)
    logInfo(log"Deleting ${MDC(LogKeys.NUM_FILES, filesToDelete.size)} " +
      log"files not used in versions >= ${MDC(LogKeys.VERSION_NUM, minVersionToRetain)}")
    var failedToDelete = 0
    filesToDelete.foreach { case (dfsFileName, maxUsedVersion) =>
      try {
        val dfsFile = dfsFilePath(dfsFileName)
        fm.delete(dfsFile)
        if (maxUsedVersion == -1) {
          logDebug(s"Deleted orphan file $dfsFileName")
        } else {
          logDebug(s"Deleted file $dfsFileName that was last used in version $maxUsedVersion")
        }
      } catch {
        case e: Exception =>
          failedToDelete += 1
          if (maxUsedVersion == -1) {
            logWarning(log"Error deleting orphan file ${MDC(LogKeys.PATH, dfsFileName)}", e)
          } else {
            logWarning(log"Error deleting file ${MDC(LogKeys.PATH, dfsFileName)}, " +
              log"last used in version ${MDC(LogKeys.MAX_FILE_VERSION, maxUsedVersion)}", e)
          }
      }
    }

    // Delete the version files and forget about them
    snapshotVersionsAndUniqueIdsToDelete.foreach { case (version, uniqueId) =>
      val versionFile = dfsBatchZipFile(version, uniqueId)
      try {
        fm.delete(versionFile)
        versionToRocksDBFiles.remove((version, uniqueId))
        logDebug(s"Deleted version $version")
      } catch {
        case e: Exception =>
          logWarning(log"Error deleting version file ${MDC(LogKeys.PATH, versionFile)} for " +
            log"version ${MDC(LogKeys.FILE_VERSION, version)}", e)
      }
    }
    logInfo(log"Deleted ${MDC(LogKeys.NUM_FILES, filesToDelete.size - failedToDelete)} files " +
      log"(failed to delete" +
      log"${MDC(LogKeys.NUM_FILES_FAILED_TO_DELETE, failedToDelete)} files) " +
      log"not used in versions >= ${MDC(LogKeys.MIN_VERSION_NUM, minVersionToRetain)}")

    val changelogVersionsAndUniqueIdsToDelete: Array[(Long, Option[String])] = changelogFiles
      .map(_.getName.stripSuffix(".changelog").split("_"))
      .map {
        case Array(version, uniqueId) => (version.toLong, Option(uniqueId))
        case Array(version) => (version.toLong, None)
      }
      .filter(_._1 < minVersionToRetain)

    deleteChangelogFiles(changelogVersionsAndUniqueIdsToDelete)

    // Delete orphan checksum files
    val checksumFiles = allFiles
      .filter(ChecksumCheckpointFileManager.isChecksumFile)
      .map(ChecksumFile)

    // Do this only if we see checksum files in the initial dir listing
    // To avoid checking for orphan checksum files, if there is no checksum files
    // (e.g. file checksum was never enabled)
    if (checksumFiles.nonEmpty) {
      // Set of changelog and zip files that were deleted.
      // They were deleted using fm, hence their checksum files were deleted too.
      val deletedStoreFiles =
        changelogVersionsAndUniqueIdsToDelete.map { case (v, id) => dfsChangelogFile(v, id) } ++
          snapshotVersionsAndUniqueIdsToDelete.map { case (v, id) => dfsBatchZipFile(v, id) }

      StateStoreProvider.deleteOrphanChecksumFiles(
        fm,
        checksumFiles.toSeq,
        deletedStoreFiles.toSeq,
        minVersionToRetain,
        fileChecksumEnabled
      )
    }

    // Always set minSeenVersion for regular deletion frequency even if deletion fails.
    // This is safe because subsequent calls retry deleting old version files
    minSeenVersion = minVersionToRetain
  }

  /** Save immutable files to DFS directory */
  private def saveImmutableFilesToDfs(
      version: Long,
      localFiles: Seq[File],
      fileMappings: Map[String, RocksDBSnapshotFile],
      checkpointUniqueId: Option[String] = None): Seq[RocksDBImmutableFile] = {
    // Get the immutable files used in previous versions, as some of those uploaded files can be
    // reused for this version
    logInfo(log"Saving RocksDB files to DFS for version ${MDC(LogKeys.VERSION_NUM, version)} " +
      log"uniqueId: ${MDC(LogKeys.UUID, checkpointUniqueId.getOrElse(""))}")

    var bytesCopied = 0L
    var filesCopied = 0L
    var filesReused = 0L

    val immutableFiles = localFiles.map { localFile =>
      val dfsFileMapping = fileMappings.get(localFile.getName)
      assert(dfsFileMapping.isDefined)
      val dfsFile = dfsFileMapping.get.immutableFile
      val existsInDfs = dfsFileMapping.get.isUploaded

      if (existsInDfs) {
        logInfo(log"reusing file ${MDC(LogKeys.DFS_FILE, dfsFile)} for " +
          log"${MDC(LogKeys.FILE_NAME, localFile)}")
        filesReused += 1
      } else {
        // Note: The implementation of copyFromLocalFile() closes the output stream when there is
        // any exception while copying. So this may generate partial files on DFS. But that is
        // okay because until the main [version].zip file is written, those partial files are
        // not going to be used at all. Eventually these files should get cleared.
        fs.copyFromLocalFile(
          new Path(localFile.getAbsoluteFile.toURI), dfsFilePath(dfsFile.dfsFileName))
        val localFileSize = localFile.length()
        logInfo(log"Copied ${MDC(LogKeys.FILE_NAME, localFile)} to " +
          log"${MDC(LogKeys.DFS_FILE, dfsFile)} - ${MDC(LogKeys.NUM_BYTES, localFileSize)} bytes")
        filesCopied += 1
        bytesCopied += localFileSize
      }

      dfsFile
    }
    logInfo(log"Copied ${MDC(LogKeys.NUM_FILES_COPIED, filesCopied)} files " +
      log"(${MDC(LogKeys.NUM_BYTES, bytesCopied)} bytes) from local to" +
      log" DFS for version ${MDC(LogKeys.VERSION_NUM, version)}. " +
      log"${MDC(LogKeys.NUM_FILES_REUSED, filesReused)} files reused without copying.")
    versionToRocksDBFiles.put((version, checkpointUniqueId), immutableFiles)
    saveCheckpointMetrics = RocksDBFileManagerMetrics(
      bytesCopied = bytesCopied,
      filesCopied = filesCopied,
      filesReused = filesReused)

    immutableFiles
  }

  /**
   * Copy files from DFS directory to a local directory. It will figure out which
   * existing files are needed, and accordingly, unnecessary SST files are deleted while
   * necessary and non-existing files are copied from DFS.
   */
  private def loadImmutableFilesFromDfs(
      immutableFiles: Seq[RocksDBImmutableFile],
      localDir: File,
      rocksDBFileMapping: RocksDBFileMapping,
      version: Long): Unit = {
    val requiredFileNameToFileDetails = immutableFiles.map(f => f.localFileName -> f).toMap

    val localImmutableFiles = listRocksDBFiles(localDir)._1

    // Delete unnecessary local immutable files
    localImmutableFiles.foreach { existingFile =>
      val existingFileSize = existingFile.length()
      val requiredFile = requiredFileNameToFileDetails.get(existingFile.getName)
      val prevDfsFile = rocksDBFileMapping.getDfsFileForLoad(this, existingFile.getName, version)
      val isSameFile = if (requiredFile.isDefined && prevDfsFile.isDefined) {
        requiredFile.get.dfsFileName == prevDfsFile.get.dfsFileName &&
          existingFile.length() == requiredFile.get.sizeBytes
      } else {
        false
      }

      if (!isSameFile) {
        rocksDBFileMapping.remove(existingFile.getName)
        existingFile.delete()
        logInfo(log"Deleted local file ${MDC(LogKeys.FILE_NAME, existingFile)} " +
          log"with size ${MDC(LogKeys.NUM_BYTES, existingFileSize)} mapped" +
          log" to previous dfsFile ${MDC(LogKeys.DFS_FILE, prevDfsFile.getOrElse("null"))}")
      } else {
        logInfo(log"reusing ${MDC(LogKeys.DFS_FILE, prevDfsFile)} present at " +
          log"${MDC(LogKeys.EXISTING_FILE, existingFile)} " +
          log"for ${MDC(LogKeys.FILE_NAME, requiredFile)}")
      }
    }

    // Delete remaining unnecessary local immutable file mappings.
    // Files present in the file mapping but not the filesystem may lead to
    // versionID mismatch error (SPARK-52637), so we should explicitly delete
    // them.
    rocksDBFileMapping.purgeIncompatibleMappingsForLoad(version).foreach {
      case (_, (dfsFileMappedVersion, dfsFile)) =>
        logInfo(log"Deleted local fileMapping to ${MDC(DFS_FILE, dfsFile)} because " +
          log"mapped file version ${MDC(VERSION_NUM, dfsFileMappedVersion)} was " +
          log"incompatible with versionToLoad ${MDC(VERSION_NUM, version)}")
    }

    var filesCopied = 0L
    var bytesCopied = 0L
    var filesReused = 0L
    immutableFiles.foreach { file =>
      val localFileName = file.localFileName
      val localFile = localFilePath(localDir, localFileName)
      if (!localFile.exists) {
        val dfsFile = dfsFilePath(file.dfsFileName)
        // Note: The implementation of copyToLocalFile() closes the output stream when there is
        // any exception while copying. So this may generate partial files on DFS. But that is
        // okay because until the main [version].zip file is written, those partial files are
        // not going to be used at all. Eventually these files should get cleared.
        fs.copyToLocalFile(dfsFile, new Path(localFile.getAbsoluteFile.toURI))
        val localFileSize = localFile.length()
        val expectedSize = file.sizeBytes
        if (localFileSize != expectedSize) {
          throw QueryExecutionErrors.unexpectedFileSize(dfsFile, localFile, expectedSize,
            localFileSize)
        }
        filesCopied += 1
        bytesCopied += localFileSize
        rocksDBFileMapping.mapToDfsFile(localFileName, file, version)
        logInfo(log"Copied ${MDC(LogKeys.DFS_FILE, dfsFile)} to " +
          log"${MDC(LogKeys.FILE_NAME, localFile)} - " +
          log"${MDC(LogKeys.NUM_BYTES, localFileSize)} bytes")
      } else {
        filesReused += 1
      }
    }
    logInfo(log"Copied ${MDC(LogKeys.NUM_FILES_COPIED, filesCopied)} files " +
      log"(${MDC(LogKeys.NUM_BYTES, bytesCopied)} bytes) from DFS to local with " +
      log"${MDC(LogKeys.NUM_FILES_REUSED, filesReused)} files reused.")

    loadCheckpointMetrics = RocksDBFileManagerMetrics(
      bytesCopied = bytesCopied,
      filesCopied = filesCopied,
      filesReused = filesReused)
  }

  /**
   * Unzip the batch zip file for the given version from DFS into localDir.
   * Zip files with the same content can have different bytes hence different checksums.
   * Hence for v1, since the zip file name is fixed (e.g. "5.zip"), we want to avoid
   * comparing with the wrong checksum file, in a situation such as concurrent uploads.
   */
  private def unzipBatchZipFileFromDfs(
      version: Long, checkpointUniqueId: Option[String], localDir: File): Seq[File] = {
    if (checkpointUniqueId.isDefined) {
      Utils.unzipFilesFromInputStream(
        fm.open(dfsBatchZipFile(version, checkpointUniqueId)), localDir)
    } else {
      Utils.unzipFilesFromFile(fs, dfsBatchZipFile(version, checkpointUniqueId), localDir)
    }
  }

  /** Get the SST files required for a version from the version zip file in DFS */
  private def getImmutableFilesFromVersionZip(
      version: Long, checkpointUniqueId: Option[String] = None): Seq[RocksDBImmutableFile] = {
    Utils.deleteRecursively(localTempDir)
    Utils.createDirectory(localTempDir)
    unzipBatchZipFileFromDfs(version, checkpointUniqueId, localTempDir)
    val metadataFile = localMetadataFile(localTempDir)
    val metadata = RocksDBCheckpointMetadata.readFromFile(metadataFile)
    metadata.immutableFiles
  }

  /**
   * Compress files to a single zip file in DFS. Only the file names are embedded in the zip.
   * Any error while writing will ensure that the file is not written.
   */
  private def zipToDfsFile(
      files: Seq[File],
      dfsZipFile: Path,
      verifyNonEmptyFilesInZip: Boolean): Unit = {
    if (verifyNonEmptyFilesInZip) {
      // Verify that all files are non-empty
      files.foreach { file =>
        // We can have an empty log file, even when WAL is disabled
        if (!isLogFile(file.getName) && file.length() == 0) {
          throw StateStoreErrors.unexpectedEmptyFileInRocksDBZip(file.getName, dfsZipFile.toString)
        }
      }
    }

    lazy val filesStr = s"$dfsZipFile\n\t${files.mkString("\n\t")}"
    var in: InputStream = null
    val out = fm.createAtomic(dfsZipFile, overwriteIfPossible = true)
    var totalBytes = 0L
    val zout = new ZipOutputStream(out)
    try {
      files.foreach { file =>
        zout.putNextEntry(new ZipEntry(file.getName))
        in = new FileInputStream(file)
        val bytes = in.transferTo(zout)
        in.close()
        zout.closeEntry()
        totalBytes += bytes
      }
      zout.close()  // so that any error in closing also cancels the output stream
      logDebug(log"Zipped ${MDC(LogKeys.NUM_BYTES, totalBytes)} bytes (before compression) to " +
        log"${MDC(LogKeys.FILE_NAME, filesStr)}")
      // The other fields saveCheckpointMetrics should have been filled
      saveCheckpointMetrics =
        saveCheckpointMetrics.copy(zipFileBytesUncompressed = Some(totalBytes))
    } catch {
      case e: Exception =>
        // Cancel the actual output stream first, so that zout.close() does not write the file
        out.cancel()
        logError(log"Error zipping to ${MDC(LogKeys.FILE_NAME, filesStr)}", e)
        throw e
    } finally {
      // Close everything no matter what happened
      Utils.closeQuietly(in)
      Utils.closeQuietly(zout)
    }
  }

  /** Log the files present in a directory. This is useful for debugging. */
  private def logFilesInDir(dir: File, msg: MessageWithContext): Unit = {
    lazy val files = Option(Utils.recursiveList(dir)).getOrElse(Array.empty).map { f =>
      s"${f.getAbsolutePath} - ${f.length()} bytes"
    }
    logDebug(msg + log" - ${MDC(LogKeys.NUM_FILES, files.length)} files\n\t" +
      log"${MDC(LogKeys.FILE_NAME, files.mkString("\n\t"))}")
  }

  def newDFSFileName(localFileName: String, dfsFileSuffix: String): String = {
    val baseName = FilenameUtils.getBaseName(localFileName)
    val extension = FilenameUtils.getExtension(localFileName)
    s"$baseName-$dfsFileSuffix.$extension"
  }

  def dfsFileSuffix(immutableFile: RocksDBImmutableFile): String = {
    val suffixStart = immutableFile.dfsFileName.indexOf('-')
    val suffixEnd = immutableFile.dfsFileName.indexOf('.')

    immutableFile.dfsFileName.substring(suffixStart + 1, suffixEnd)
  }

  private def dfsBatchZipFile(version: Long, checkpointUniqueId: Option[String] = None): Path =
    checkpointUniqueId.map(id => new Path(s"$dfsRootDir/${version}_$id.zip"))
      .getOrElse(new Path(s"$dfsRootDir/$version.zip"))
  // We use changelog suffix intentionally so that we can tell the difference from changelog file of
  // HDFSBackedStateStore which is named version.delta.
  private def dfsChangelogFile(version: Long, checkpointUniqueId: Option[String] = None): Path =
    checkpointUniqueId.map(id => new Path(s"$dfsRootDir/${version}_$id.changelog"))
      .getOrElse(new Path(s"$dfsRootDir/$version.changelog"))

  private def localMetadataFile(parentDir: File): File = new File(parentDir, "metadata")

  override protected def logName: String = s"${super.logName} $loggingId"

  private def dfsFilePath(fileName: String): Path = {
    if (isSstFile(fileName)) {
      new Path(new Path(dfsRootDir, SST_FILES_DFS_SUBDIR), fileName)
    } else if (isLogFile(fileName)) {
      new Path(new Path(dfsRootDir, LOG_FILES_DFS_SUBDIR), fileName)
    } else {
      new Path(dfsRootDir, fileName)
    }
  }

  private def localFilePath(localDir: File, fileName: String): File = {
    if (isLogFile(fileName)) {
      new File(new File(localDir, LOG_FILES_LOCAL_SUBDIR), fileName)
    } else {
      new File(localDir, fileName)
    }
  }

  /**
   * List all the RocksDB files that need be synced or recovered.
   */
  def listRocksDBFiles(localDir: File): (Seq[File], Seq[File]) = {
    val topLevelFiles = localDir.listFiles.filter(!_.isDirectory)
    val archivedLogFiles =
      Option(new File(localDir, LOG_FILES_LOCAL_SUBDIR).listFiles())
        .getOrElse(Array[File]())
        // To ignore .log.crc files
        .filter(file => isLogFile(file.getName))
    val (topLevelSstFiles, topLevelOtherFiles) = topLevelFiles.partition(f => isSstFile(f.getName))
    ((topLevelSstFiles ++ archivedLogFiles).toImmutableArraySeq,
      topLevelOtherFiles.toImmutableArraySeq)
  }
}

/**
 * Metrics regarding RocksDB file sync between local and DFS.
 */
case class RocksDBFileManagerMetrics(
    filesCopied: Long,
    bytesCopied: Long,
    filesReused: Long,
    @JsonDeserialize(contentAs = classOf[java.lang.Long])
    zipFileBytesUncompressed: Option[Long] = None,
    @JsonDeserialize(contentAs = classOf[java.lang.Long])
    saveZipFilesTimeMs: Option[Long] = None)

/**
 * Metrics to return when requested but no operation has been performed.
 */
object RocksDBFileManagerMetrics {
  val EMPTY_METRICS = RocksDBFileManagerMetrics(0L, 0L, 0L, None)
}

/**
 * Case class to keep track of column family info within checkpoint metadata.
 * @param cfId - virtual column family id
 * @param isInternal - whether the column family is internal or not
 */
case class ColumnFamilyInfo(
    cfId: Short,
    isInternal: Boolean)

/**
 * Classes to represent metadata of checkpoints saved to DFS. Since this is converted to JSON, any
 * changes to this MUST be backward-compatible.
 */
case class RocksDBCheckpointMetadata(
    sstFiles: Seq[RocksDBSstFile],
    logFiles: Seq[RocksDBLogFile],
    numKeys: Long,
    numInternalKeys: Long,
    columnFamilyMapping: Option[Map[String, Short]] = None,
    columnFamilyTypeMap: Option[Map[String, Boolean]] = None,
    maxColumnFamilyId: Option[Short] = None) {

  require(columnFamilyMapping.isDefined == maxColumnFamilyId.isDefined,
    "columnFamilyMapping and maxColumnFamilyId must either both be defined or both be None")

  import RocksDBCheckpointMetadata._

  def json: String = {
    // We turn this field into a null to avoid write an empty logFiles field in the json.
    val nullified = if (logFiles.isEmpty) this.copy(logFiles = null) else this
    mapper.writeValueAsString(nullified)
  }

  def prettyJson: String = Serialization.writePretty(this)(RocksDBCheckpointMetadata.format)

  def writeToFile(metadataFile: File): Unit = {
    val writer = Files.newBufferedWriter(metadataFile.toPath, UTF_8)
    try {
      writer.write(s"v$VERSION\n")
      writer.write(this.json)
    } finally {
      writer.close()
    }
  }

  def immutableFiles: Seq[RocksDBImmutableFile] = sstFiles ++ logFiles
}

/** Helper class for [[RocksDBCheckpointMetadata]] */
object RocksDBCheckpointMetadata {
  val VERSION = 1

  implicit val format: Formats = Serialization.formats(NoTypeHints)

  /** Used to convert between classes and JSON. */
  lazy val mapper = {
    val _mapper = new ObjectMapper with ClassTagExtensions
    _mapper.setDefaultPropertyInclusion(Include.NON_ABSENT)
    _mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false)
    _mapper.registerModule(DefaultScalaModule)
    _mapper
  }

  def readFromFile(metadataFile: File): RocksDBCheckpointMetadata = {
    val reader = Files.newBufferedReader(metadataFile.toPath, UTF_8)
    try {
      val versionLine = reader.readLine()
      if (versionLine != s"v$VERSION") {
        throw QueryExecutionErrors.cannotReadCheckpoint(s"v$VERSION", versionLine)
      }
      Serialization.read[RocksDBCheckpointMetadata](reader)
    } finally {
      reader.close()
    }
  }

  // Apply method for cases without column family information
  def apply(
      rocksDBFiles: Seq[RocksDBImmutableFile],
      numKeys: Long): RocksDBCheckpointMetadata = {
    val (sstFiles, logFiles) = rocksDBFiles.partition(_.isInstanceOf[RocksDBSstFile])
    new RocksDBCheckpointMetadata(
      sstFiles.map(_.asInstanceOf[RocksDBSstFile]),
      logFiles.map(_.asInstanceOf[RocksDBLogFile]),
      numKeys,
      0,
      None,
      None
    )
  }

  def apply(
      rocksDBFiles: Seq[RocksDBImmutableFile],
      numKeys: Long,
      numInternalKeys: Long,
      columnFamilyMapping: Option[Map[String, Short]],
      columnFamilyTypeMap: Option[Map[String, Boolean]],
      maxColumnFamilyId: Option[Short]): RocksDBCheckpointMetadata = {
    val (sstFiles, logFiles) = rocksDBFiles.partition(_.isInstanceOf[RocksDBSstFile])
    new RocksDBCheckpointMetadata(
      sstFiles.map(_.asInstanceOf[RocksDBSstFile]),
      logFiles.map(_.asInstanceOf[RocksDBLogFile]),
      numKeys,
      numInternalKeys,
      columnFamilyMapping,
      columnFamilyTypeMap,
      maxColumnFamilyId
    )
  }

  // Apply method for cases with separate sstFiles and logFiles, without column family information
  def apply(
      sstFiles: Seq[RocksDBSstFile],
      logFiles: Seq[RocksDBLogFile],
      numKeys: Long): RocksDBCheckpointMetadata = {
    new RocksDBCheckpointMetadata(sstFiles, logFiles, numKeys, 0, None, None)
  }

  // Apply method for cases with column family information
  def apply(
      rocksDBFiles: Seq[RocksDBImmutableFile],
      numKeys: Long,
      numInternalKeys: Long,
      columnFamilyMapping: Map[String, Short],
      columnFamilyTypeMap: Map[String, Boolean],
      maxColumnFamilyId: Short): RocksDBCheckpointMetadata = {
    val (sstFiles, logFiles) = rocksDBFiles.partition(_.isInstanceOf[RocksDBSstFile])
    new RocksDBCheckpointMetadata(
      sstFiles.map(_.asInstanceOf[RocksDBSstFile]),
      logFiles.map(_.asInstanceOf[RocksDBLogFile]),
      numKeys,
      numInternalKeys,
      Some(columnFamilyMapping),
      Some(columnFamilyTypeMap),
      Some(maxColumnFamilyId)
    )
  }

  // Apply method for cases with separate sstFiles and logFiles, and column family information
  def apply(
      sstFiles: Seq[RocksDBSstFile],
      logFiles: Seq[RocksDBLogFile],
      numKeys: Long,
      numInternalKeys: Long,
      columnFamilyMapping: Map[String, Short],
      columnFamilyTypeMap: Map[String, Boolean],
      maxColumnFamilyId: Short): RocksDBCheckpointMetadata = {
    new RocksDBCheckpointMetadata(
      sstFiles,
      logFiles,
      numKeys,
      numInternalKeys,
      Some(columnFamilyMapping),
      Some(columnFamilyTypeMap),
      Some(maxColumnFamilyId)
    )
  }
}

/**
 * A RocksDBImmutableFile maintains a mapping between a local RocksDB file name and the name of
 * its copy on DFS. Since these files are immutable, their DFS copies can be reused.
 */
sealed trait RocksDBImmutableFile {
  def localFileName: String
  def dfsFileName: String
  def sizeBytes: Long

  /**
   * Whether another local file is same as the file described by this class.
   * A file is same only when the name and the size are same.
   */
  def isSameFile(otherFile: File): Boolean = {
    otherFile.getName == localFileName && otherFile.length() == sizeBytes
  }
}

/**
 * Class to represent a RocksDB SST file. Since this is converted to JSON,
 * any changes to these MUST be backward-compatible.
 */
private[sql] case class RocksDBSstFile(
    localFileName: String,
    dfsSstFileName: String,
    sizeBytes: Long) extends RocksDBImmutableFile {

  override def dfsFileName: String = dfsSstFileName
}

/**
 * Class to represent a RocksDB Log file. Since this is converted to JSON,
 * any changes to these MUST be backward-compatible.
 */
private[sql] case class RocksDBLogFile(
    localFileName: String,
    dfsLogFileName: String,
    sizeBytes: Long) extends RocksDBImmutableFile {

  override def dfsFileName: String = dfsLogFileName
}

object RocksDBImmutableFile {
  val SST_FILES_DFS_SUBDIR = "SSTs"
  val LOG_FILES_DFS_SUBDIR = "logs"
  val LOG_FILES_LOCAL_SUBDIR = "archive"

  def apply(
      localFileName: String,
      dfsFileName: String,
      sizeBytes: Long): RocksDBImmutableFile = {
    if (isSstFile(localFileName)) {
      RocksDBSstFile(localFileName, dfsFileName, sizeBytes)
    } else if (isLogFile(localFileName)) {
      RocksDBLogFile(localFileName, dfsFileName, sizeBytes)
    } else {
      null
    }
  }

  def isSstFile(fileName: String): Boolean = fileName.endsWith(".sst")

  def isLogFile(fileName: String): Boolean = fileName.endsWith(".log")

  private def isArchivedLogFile(file: File): Boolean =
    isLogFile(file.getName) && file.getParentFile.getName == LOG_FILES_LOCAL_SUBDIR

  def isImmutableFile(file: File): Boolean = isSstFile(file.getName) || isArchivedLogFile(file)
}
