/*
 * 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.paimon.spark.procedure

import org.apache.paimon.{utils, Snapshot}
import org.apache.paimon.catalog.{Catalog, Identifier}
import org.apache.paimon.fs.Path
import org.apache.paimon.manifest.{ManifestEntry, ManifestFile}
import org.apache.paimon.operation.{CleanOrphanFilesResult, OrphanFilesClean}
import org.apache.paimon.operation.OrphanFilesClean.retryReadingFiles
import org.apache.paimon.table.FileStoreTable
import org.apache.paimon.utils.FileStorePathFactory.BUCKET_PATH_PREFIX
import org.apache.paimon.utils.SerializableConsumer

import org.apache.spark.internal.Logging
import org.apache.spark.sql.{functions, Dataset, PaimonSparkSession, SparkSession}
import org.apache.spark.sql.catalyst.SQLConfHelper

import java.util
import java.util.Collections
import java.util.concurrent.atomic.AtomicLong
import java.util.function.Consumer

import scala.collection.JavaConverters._
import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer

case class SparkOrphanFilesClean(
    specifiedTable: FileStoreTable,
    specifiedOlderThanMillis: Long,
    parallelism: Int,
    dryRunPara: Boolean,
    @transient spark: SparkSession)
  extends OrphanFilesClean(specifiedTable, specifiedOlderThanMillis, dryRunPara)
  with SQLConfHelper
  with Logging {

  def doOrphanClean(): (Dataset[(Long, Long)], Dataset[BranchAndManifestFile]) = {
    import spark.implicits._

    val branches = validBranches()
    val deletedFilesCountInLocal = new AtomicLong(0)
    val deletedFilesLenInBytesInLocal = new AtomicLong(0)
    // snapshot and changelog files are the root of everything, so they are handled specially
    // here, and subsequently, we will not count their orphan files.
    cleanSnapshotDir(
      branches,
      (_: Path) => deletedFilesCountInLocal.incrementAndGet,
      size => deletedFilesLenInBytesInLocal.addAndGet(size))

    val maxBranchParallelism = Math.min(branches.size(), parallelism)
    // find snapshots using branch and find manifests(manifest, index, statistics) using snapshot
    val usedManifestFiles = spark.sparkContext
      .parallelize(branches.asScala.toSeq, maxBranchParallelism)
      .mapPartitions(_.flatMap {
        branch => safelyGetAllSnapshots(branch).asScala.map(snapshot => (branch, snapshot.toJson))
      })
      .repartition(parallelism)
      .flatMap {
        case (branch, snapshotJson) =>
          val usedFileBuffer = new ArrayBuffer[BranchAndManifestFile]()
          val usedFileConsumer =
            new Consumer[org.apache.paimon.utils.Pair[String, java.lang.Boolean]] {
              override def accept(pair: utils.Pair[String, java.lang.Boolean]): Unit = {
                usedFileBuffer.append(BranchAndManifestFile(branch, pair.getLeft, pair.getRight))
              }
            }
          val snapshot = Snapshot.fromJson(snapshotJson)
          collectWithoutDataFileWithManifestFlag(branch, snapshot, usedFileConsumer)
          usedFileBuffer
      }
      .toDS()
      .cache()

    // find all data files
    val dataFiles = usedManifestFiles
      .filter(_.isManifestFile)
      .distinct()
      .mapPartitions {
        it =>
          val branchManifests = new util.HashMap[String, ManifestFile]
          it.flatMap {
            branchAndManifestFile =>
              val manifestFile = branchManifests.computeIfAbsent(
                branchAndManifestFile.branch,
                (key: String) =>
                  specifiedTable.switchToBranch(key).store.manifestFileFactory.create)

              retryReadingFiles(
                () => manifestFile.readWithIOException(branchAndManifestFile.manifestName),
                Collections.emptyList[ManifestEntry]
              ).asScala.flatMap {
                manifestEntry =>
                  manifestEntry.fileName() +: manifestEntry.file().extraFiles().asScala
              }
          }
      }

    // union manifest and data files
    val usedFiles = usedManifestFiles
      .map(_.manifestName)
      .union(dataFiles)
      .toDF("used_name")

    // find candidate files which can be removed
    val fileDirs = listPaimonFileDirs.asScala.map(_.toString).toSeq
    val maxFileDirsParallelism = Math.min(fileDirs.size, parallelism)
    val candidates = spark.sparkContext
      .parallelize(fileDirs, maxFileDirsParallelism)
      .flatMap {
        dir =>
          tryBestListingDirs(new Path(dir)).asScala
            .filter(file => !file.isDir())
            .filter(file => !isManagedBlobPack(file.getPath))
            .filter(oldEnough)
            .map {
              file =>
                val path = file.getPath
                (path.getName, path.toString, file.getLen, path.getParent.toString)
            }
      }
      .toDF("name", "path", "len", "dataDir")
      .repartition(parallelism)

    // use left anti to filter files which is not used
    val deleted = candidates
      .join(usedFiles, $"name" === $"used_name", "left_anti")
      .repartition($"dataDir")
      .mapPartitions {
        it =>
          var deletedFilesCount = 0L
          var deletedFilesLenInBytes = 0L
          val dataDirs = new mutable.HashSet[String]()

          while (it.hasNext) {
            val fileInfo = it.next();
            val pathToClean = fileInfo.getString(1)
            val deletedPath = new Path(pathToClean)
            deletedFilesLenInBytes += fileInfo.getLong(2)
            cleanFile(deletedPath)
            logInfo(s"Cleaned file: $pathToClean")
            dataDirs.add(fileInfo.getString(3))
            deletedFilesCount += 1
          }

          // clean empty directory
          if (!dryRun) {
            val bucketDirs = dataDirs
              .filter(_.contains(BUCKET_PATH_PREFIX))
              .map(new Path(_))
            tryCleanDataDirectory(bucketDirs.asJava, partitionKeysNum + 1)
          }

          logInfo(
            s"Total cleaned files: $deletedFilesCount, Total cleaned files len : $deletedFilesLenInBytes")
          Iterator.single((deletedFilesCount, deletedFilesLenInBytes))
      }

    val finalDeletedDataset =
      if (deletedFilesCountInLocal.get() != 0 || deletedFilesLenInBytesInLocal.get() != 0) {
        deleted.union(
          spark.createDataset(
            Seq((deletedFilesCountInLocal.get(), deletedFilesLenInBytesInLocal.get()))))
      } else {
        deleted
      }

    (finalDeletedDataset, usedManifestFiles)
  }
}

/**
 * @param branch
 *   The branch name
 * @param manifestName
 *   The manifest file name, including manifest-list, manifest, index-manifest, statistics
 * @param isManifestFile
 *   If it is the manifest file
 */
case class BranchAndManifestFile(branch: String, manifestName: String, isManifestFile: Boolean)

object SparkOrphanFilesClean extends SQLConfHelper {
  def executeDatabaseOrphanFiles(
      catalog: Catalog,
      databaseName: String,
      tableName: String,
      olderThanMillis: Long,
      parallelismOpt: Integer,
      dryRun: Boolean): CleanOrphanFilesResult = {
    val spark = PaimonSparkSession.active
    val parallelism = if (parallelismOpt == null) {
      Math.max(spark.sparkContext.defaultParallelism, conf.numShufflePartitions)
    } else {
      parallelismOpt.intValue()
    }

    val tableNames = if (tableName == null || "*" == tableName) {
      catalog.listTables(databaseName).asScala
    } else {
      tableName :: Nil
    }
    val tables = tableNames.map {
      tableName =>
        val identifier = new Identifier(databaseName, tableName)
        val table = catalog.getTable(identifier)
        assert(
          table.isInstanceOf[FileStoreTable],
          s"Only FileStoreTable supports remove-orphan-files action. The table type is '${table.getClass.getName}'.")
        table.asInstanceOf[FileStoreTable]
    }
    if (tables.isEmpty) {
      return new CleanOrphanFilesResult(0, 0)
    }
    val (deleted, waitToRelease) = tables.map {
      table =>
        new SparkOrphanFilesClean(
          table,
          olderThanMillis,
          parallelism,
          dryRun,
          spark
        ).doOrphanClean()
    }.unzip
    try {
      val result = deleted
        .reduce((l, r) => l.union(r))
        .toDF("deletedFilesCount", "deletedFilesLenInBytes")
        .agg(functions.sum("deletedFilesCount"), functions.sum("deletedFilesLenInBytes"))
        .head()
      assert(result.schema.size == 2, result.schema)
      if (result.isNullAt(0)) {
        // no files can be deleted
        new CleanOrphanFilesResult(0, 0)
      } else {
        new CleanOrphanFilesResult(result.getLong(0), result.getLong(1))
      }
    } finally {
      waitToRelease.foreach(_.unpersist())
    }
  }
}
