/*
 * 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

import org.apache.paimon.spark.read.ObjectTableScanBuilder
import org.apache.paimon.spark.rowops.{PaimonSparkCopyOnWriteOperation, PaimonSparkDeltaOperation}
import org.apache.paimon.table.`object`.ObjectTable
import org.apache.paimon.table.{BucketMode, FileStoreTable, Table}

import org.apache.spark.sql.connector.catalog.{SupportsRead, SupportsRowLevelOperations, TableCapability}
import org.apache.spark.sql.connector.read.ScanBuilder
import org.apache.spark.sql.connector.write.{RowLevelOperationBuilder, RowLevelOperationInfo}
import org.apache.spark.sql.util.CaseInsensitiveStringMap

import java.util.{EnumSet => JEnumSet, Set => JSet}

/**
 * A spark [[org.apache.spark.sql.connector.catalog.Table]] for paimon.
 *
 * This base variant does NOT implement [[SupportsRowLevelOperations]]. Spark 4.1 moved
 * `RewriteDeleteFromTable` / `RewriteUpdateTable` / `RewriteMergeIntoTable` from the separate "DML
 * rewrite" batch into the main Resolution batch, which fires BEFORE Paimon's own
 * postHocResolutionRule interceptors (`PaimonDeleteTable`, `PaimonUpdateTable`, `PaimonMergeInto`).
 * If this base class implemented `SupportsRowLevelOperations`, Spark 4.1 would immediately call
 * `newRowLevelOperationBuilder` on tables whose V2 write is disabled (e.g. dynamic bucket or
 * primary-key tables that fall back to V1 write) and fail before Paimon has a chance to rewrite the
 * plan to a V1 command. Likewise, data-evolution and fixed-length CHAR tables need to stay on
 * Paimon's V1 postHoc path even when `useV2Write=true`, so they must also not expose
 * `SupportsRowLevelOperations`.
 *
 * Tables that DO support V2 row-level operations use the [[SparkTableWithRowLevelOps]] subclass
 * instead; the [[SparkTable.of]] factory picks the right variant via
 * [[SparkTable.supportsV2RowLevelOps]]. Append-only tables, including row-tracking-only tables,
 * expose `SupportsRowLevelOperations` so DELETE, UPDATE, and MERGE INTO can go through the V2
 * copy-on-write path when the table has no PK, deletion vectors, data evolution, or CHAR columns;
 * unaware-bucket deletion-vector append tables expose it so DELETE can go through the delta path
 * (see [[SparkTable.supportsV2DeltaOps]]).
 */
case class SparkTable(override val table: Table) extends PaimonSparkTableBase(table)

/**
 * A Paimon [[SparkTable]] that additionally exposes V2 row-level operations to Spark. Constructed
 * when [[PaimonSparkTableBase.useV2Write]] is true.
 */
class SparkTableWithRowLevelOps(tableArg: Table)
  extends SparkTable(tableArg)
  with SupportsRowLevelOperations {

  override def newRowLevelOperationBuilder(
      info: RowLevelOperationInfo): RowLevelOperationBuilder = {
    table match {
      case t: FileStoreTable =>
        if (SparkTable.supportsV2DeltaOps(this)) { () => new PaimonSparkDeltaOperation(t, info) }
        else { () => new PaimonSparkCopyOnWriteOperation(t, info) }
      case _ =>
        throw new UnsupportedOperationException(
          "Row-level write operation is only supported for FileStoreTable. " +
            s"Actual table type: ${table.getClass.getSimpleName}")
    }
  }
}

/** Factory helpers for constructing the right [[SparkTable]] subclass. */
object SparkTable {

  /**
   * Returns a [[SparkTable]] variant suitable for the given table: [[SparkTableWithRowLevelOps]]
   * when the table can participate in Spark's V2 row-level ops path, the plain [[SparkTable]]
   * otherwise.
   */
  def of(table: Table): SparkTable = {
    // We need `useV2Write` (and other coreOptions) which are instance state, so construct once to
    // probe and promote to the row-level-ops variant only if the table truly supports the V2
    // row-level write path. The base instance is cheap and is discarded if we decide to return
    // the subclass.
    val base = SparkTable(table)
    if (supportsV2RowLevelOps(base)) new SparkTableWithRowLevelOps(table) else base
  }

  /**
   * Whether the given table supports Paimon's V2 row-level operations, i.e. whether it is safe to
   * expose [[SupportsRowLevelOperations]] to Spark.
   *
   * Append-only tables return `true` here so that `SparkTable.of` wraps them as
   * `SparkTableWithRowLevelOps`, enabling Spark's V2 copy-on-write DELETE, UPDATE, and MERGE INTO
   * paths. Row-tracking append-only tables require Spark 4.0+ because Spark 3.5 does not have the
   * metadata-aware `DataWriter.write(metadata, data)` path needed to preserve row-tracking metadata
   * for rewritten rows.
   *
   * Per-version shims for Spark 3.2/3.3/3.4 each ship their own
   * `org.apache.paimon.spark.SparkTable` (class + companion) that shadows this one at packaging
   * time — the common jar is also shaded into every per-version artifact, and the per-version copy
   * wins classloader precedence. Those shim companions MUST therefore expose a method with this
   * exact signature (they hard-code `false` because Spark < 3.5 cannot participate in V2 row-level
   * ops), otherwise `RowLevelHelper.shouldFallbackToV1` on the shim classpath throws
   * `NoSuchMethodError` at the first DML statement.
   */
  private[spark] def supportsV2RowLevelOps(sparkTable: SparkTable): Boolean = {
    supportsV2CopyOnWriteOps(sparkTable) || supportsV2DeltaOps(sparkTable)
  }

  private def supportsV2CopyOnWriteOps(sparkTable: SparkTable): Boolean = {
    if (org.apache.spark.SPARK_VERSION < "3.5") return false
    if (!sparkTable.useV2Write) return false
    sparkTable.getTable match {
      case fs: FileStoreTable =>
        val supportsRowTrackingCopyOnWrite =
          !sparkTable.coreOptions.rowTrackingEnabled() || org.apache.spark.SPARK_VERSION >= "4.0"
        fs.primaryKeys().isEmpty &&
        supportsRowTrackingCopyOnWrite &&
        !sparkTable.coreOptions.deletionVectorsEnabled() &&
        !sparkTable.coreOptions.dataEvolutionEnabled() &&
        !SparkTypeUtils.containsCharType(fs.rowType())
      case _ => false
    }
  }

  /**
   * Whether the table takes the delta-based row-level path ([[PaimonSparkDeltaOperation]]): an
   * unaware-bucket append-only table with deletion vectors enabled. Bucketed append tables are
   * excluded because commit conflict detection does not cover concurrent deletion-vector index
   * changes of the same bucket (`IndexManifestFileHandler` overwrites the whole bucket's index);
   * row-tracking deletion-vector tables are excluded until the delta path defines the row lineage
   * semantics for rewritten rows.
   *
   * The Spark 3.2/3.3/3.4 shim `SparkTable` companions must expose a method with this exact
   * signature (hard-coded `false`), same as [[supportsV2RowLevelOps]]. Public because the Spark 4.1
   * rewrite scope (`PureAppendOnlyScope.targetsV2DeltaTable`) delegates to it as the single source
   * of truth for the delta gating conditions.
   */
  def supportsV2DeltaOps(sparkTable: SparkTable): Boolean = {
    if (org.apache.spark.SPARK_VERSION < "3.5") return false
    if (!sparkTable.useV2Write) return false
    sparkTable.getTable match {
      case fs: FileStoreTable =>
        fs.primaryKeys().isEmpty &&
        fs.bucketMode() == BucketMode.BUCKET_UNAWARE &&
        sparkTable.coreOptions.deletionVectorsEnabled() &&
        !sparkTable.coreOptions.rowTrackingEnabled() &&
        !sparkTable.coreOptions.dataEvolutionEnabled() &&
        !SparkTypeUtils.containsCharType(fs.rowType())
      case _ => false
    }
  }
}

case class SparkIcebergTable(table: Table) extends BaseTable

case class SparkLanceTable(table: Table) extends BaseTable

case class SparkObjectTable(override val table: ObjectTable) extends BaseTable with SupportsRead {

  override def capabilities(): JSet[TableCapability] = {
    JEnumSet.of(TableCapability.BATCH_READ)
  }

  override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = {
    new ObjectTableScanBuilder(table.copy(options.asCaseSensitiveMap))
  }
}
