/*
 * 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.connector.catalog

import java.util
import java.util.{Objects, UUID}

import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.connector.catalog.constraints.Constraint
import org.apache.spark.sql.connector.distributions.{Distribution, Distributions}
import org.apache.spark.sql.connector.expressions.{SortOrder, Transform}
import org.apache.spark.sql.connector.read._
import org.apache.spark.sql.connector.write.{LogicalWriteInfo, SupportsOverwrite, WriteBuilder, WriterCommitMessage}
import org.apache.spark.sql.sources._
import org.apache.spark.sql.types.{LongType, StructField, StructType}
import org.apache.spark.sql.util.CaseInsensitiveStringMap
import org.apache.spark.unsafe.types.UTF8String
import org.apache.spark.util.ArrayImplicits._

/**
 * A simple in-memory table. Rows are stored as a buffered group produced by each output task.
 */
class InMemoryTable(
    name: String,
    columns: Array[Column],
    override val partitioning: Array[Transform],
    override val properties: util.Map[String, String],
    override val constraints: Array[Constraint] = Array.empty,
    distribution: Distribution = Distributions.unspecified(),
    ordering: Array[SortOrder] = Array.empty,
    numPartitions: Option[Int] = None,
    advisoryPartitionSize: Option[Long] = None,
    isDistributionStrictlyRequired: Boolean = true,
    override val numRowsPerSplit: Int = Int.MaxValue,
    override val id: String = UUID.randomUUID().toString)
  extends InMemoryBaseTable(name, columns, partitioning, properties, constraints, distribution,
    ordering, numPartitions, advisoryPartitionSize, isDistributionStrictlyRequired,
    numRowsPerSplit) with SupportsDelete {

  def this(
      name: String,
      schema: StructType,
      partitioning: Array[Transform],
      properties: util.Map[String, String]) = {
    this(
      name,
      CatalogV2Util.structTypeToV2Columns(schema),
      partitioning,
      properties)
  }

  override def canDeleteWhere(filters: Array[Filter]): Boolean = {
    InMemoryTable.supportsFilters(filters)
  }

  override def deleteWhere(filters: Array[Filter]): Unit = dataMap.synchronized {
    import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.MultipartIdentifierHelper
    dataMap --= InMemoryTable
      .filtersToKeys(dataMap.keys, partCols.map(_.toSeq.quoted).toImmutableArraySeq, filters)
    increaseVersion()
  }

  override def withData(data: Array[BufferedRows]): InMemoryTable = {
    withData(data, columns())
  }

  override def withData(data: Array[BufferedRows], columns: Array[Column]): InMemoryTable = {
    withData(data, CatalogV2Util.v2ColumnsToStructType(columns))
  }

  override def withData(
      data: Array[BufferedRows],
      writeSchema: StructType): InMemoryTable = {
    dataMap.synchronized {
      data.foreach {
        bufferedRow => {
          bufferedRow.rows.foreach { row =>
            val key = getKey(row, writeSchema)
            dataMap += dataMap.get(key)
              .map { splits =>
                val newSplits = if ((splits.last.rows.size >= numRowsPerSplit) ||
                  (splits.last.schema != writeSchema)) {
                  splits :+ new BufferedRows(key, writeSchema)
                } else {
                  splits
                }
                newSplits.last.withRow(row)
                key -> newSplits
              }
              .getOrElse(key -> Seq(new BufferedRows(key, writeSchema).withRow(row)))
            addPartitionKey(key)
          }
        }
      }

      if (data.exists(_.rows.exists(row => row.numFields == 1 &&
          row.getInt(0) == InMemoryTable.uncommittableValue()))) {
        throw new IllegalArgumentException(s"Test only mock write failure")
      }
      increaseVersion()
      this
    }
  }

  override def alterTableWithData(
      data: Array[BufferedRows],
      newSchema: StructType): InMemoryTable = {
    super.alterTableWithData(data, newSchema).asInstanceOf[InMemoryTable]
  }

  override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = {
    InMemoryBaseTable.maybeSimulateFailedTableWrite(new CaseInsensitiveStringMap(properties))
    InMemoryBaseTable.maybeSimulateFailedTableWrite(info.options)

    new InMemoryWriterBuilderWithOverWrite(info)
  }

  override def copy(): Table = {
    val copiedTable = new InMemoryTable(
      name,
      columns(),
      partitioning,
      properties,
      constraints,
      distribution,
      ordering,
      numPartitions,
      advisoryPartitionSize,
      isDistributionStrictlyRequired,
      numRowsPerSplit,
      id)

    dataMap.synchronized {
      dataMap.foreach { case (key, splits) =>
        val copiedSplits = splits.map { bufferedRows =>
          val copiedBufferedRows = new BufferedRows(bufferedRows.key, bufferedRows.schema)
          copiedBufferedRows.rows ++= bufferedRows.rows.map(_.copy())
          copiedBufferedRows
        }
        copiedTable.dataMap.put(key, copiedSplits)
      }
    }

    copiedTable.commits ++= commits.map(_.copy())

    copiedTable.setVersionAndValidatedVersionFrom(this)

    copiedTable
  }

  override def equals(other: Any): Boolean = other match {
    case that: InMemoryTable =>
      this.id == that.id && this.version() == that.version()
    case _ => false
  }

  override def hashCode(): Int = {
    Objects.hash(id, version())
  }

  class InMemoryWriterBuilderWithOverWrite(override val info: LogicalWriteInfo)
    extends InMemoryWriterBuilder(info) with SupportsOverwrite {

    override def truncate(): WriteBuilder = {
      if (!writer.isInstanceOf[Append]) {
        throw new IllegalArgumentException(s"Unsupported writer type: $writer")
      }
      writer = new TruncateAndAppend(info)
      streamingWriter = new StreamingTruncateAndAppend(info)
      this
    }

    override def overwrite(filters: Array[Filter]): WriteBuilder = {
      if (!writer.isInstanceOf[Append]) {
        throw new IllegalArgumentException(s"Unsupported writer type: $writer")
      }
      writer = new Overwrite(filters)
      streamingWriter = new StreamingNotSupportedOperation(
        s"overwrite (${filters.mkString("filters(", ", ", ")")})")
      this
    }

    override def canOverwrite(filters: Array[Filter]): Boolean = {
      InMemoryTable.supportsFilters(filters)
    }
  }

  private class Overwrite(filters: Array[Filter]) extends TestBatchWrite {
    import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.MultipartIdentifierHelper
    override protected def doCommit(
        messages: Array[WriterCommitMessage]): Unit = dataMap.synchronized {
      val deleteKeys = InMemoryTable.filtersToKeys(
        dataMap.keys, partCols.map(_.toSeq.quoted).toImmutableArraySeq, filters)
      dataMap --= deleteKeys
      withData(messages.map(_.asInstanceOf[BufferedRows]))
    }
  }
}

object InMemoryTable {

  // Convert UTF8String to string to make sure equality checks between filters and partitions
  // work correctly.
  private def valuesEqual(filterValue: Any, partitionValue: Any): Boolean =
    (filterValue, partitionValue) match {
      case (s: String, u: UTF8String) => u.toString == s
      case (u: UTF8String, s: String) => u.toString == s
      case _ => filterValue == partitionValue
    }

  def filtersToKeys(
      keys: Iterable[Seq[Any]],
      partitionNames: Seq[String],
      filters: Array[Filter]): Iterable[Seq[Any]] = {
    keys.filter { partValues =>
      filters.flatMap(splitAnd).forall {
        case EqualTo(attr, value) =>
          valuesEqual(value, InMemoryBaseTable.extractValue(attr, partitionNames, partValues))
        case EqualNullSafe(attr, value) =>
          val attrVal = InMemoryBaseTable.extractValue(attr, partitionNames, partValues)
          if (attrVal == null && value == null) {
            true
          } else if (attrVal == null || value == null) {
            false
          } else {
            valuesEqual(value, attrVal)
          }
        case IsNull(attr) =>
          null == InMemoryBaseTable.extractValue(attr, partitionNames, partValues)
        case IsNotNull(attr) =>
          null != InMemoryBaseTable.extractValue(attr, partitionNames, partValues)
        case AlwaysTrue() => true
        case f =>
          throw new IllegalArgumentException(s"Unsupported filter type: $f")
      }
    }
  }

  def supportsFilters(filters: Array[Filter]): Boolean = {
    filters.flatMap(splitAnd).forall {
      case _: EqualTo => true
      case _: EqualNullSafe => true
      case _: IsNull => true
      case _: IsNotNull => true
      case _: AlwaysTrue => true
      case _ => false
    }
  }

  def uncommittableValue(): Int = Int.MaxValue / 2

  private def splitAnd(filter: Filter): Seq[Filter] = {
    filter match {
      case And(left, right) => splitAnd(left) ++ splitAnd(right)
      case _ => filter :: Nil
    }
  }
}

/**
 * A metadata table that returns snapshot (commit) information for a parent table.
 * Simulates data source tables with multi-part identifiers, ex Iceberg's db.table.snapshots.
 */
class InMemorySnapshotsTable(parentTable: InMemoryTable) extends Table with SupportsRead {
  override def name(): String = parentTable.name + ".snapshots"

  override def schema(): StructType = StructType(Seq(
    StructField("committed_at", LongType, nullable = false),
    StructField("snapshot_id", LongType, nullable = false)
  ))

  override def capabilities(): util.Set[TableCapability] = {
    util.EnumSet.of(TableCapability.BATCH_READ)
  }

  override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = {
    new InMemorySnapshotsScanBuilder(parentTable)
  }
}

class InMemorySnapshotsScanBuilder(parentTable: InMemoryTable) extends ScanBuilder {
  override def build(): Scan = new InMemorySnapshotsScan(parentTable)
}

class InMemorySnapshotsScan(parentTable: InMemoryTable) extends Scan with Batch {
  override def readSchema(): StructType = StructType(Seq(
    StructField("committed_at", LongType, nullable = false),
    StructField("snapshot_id", LongType, nullable = false)
  ))

  override def toBatch: Batch = this

  override def planInputPartitions(): Array[InputPartition] = {
    Array(InMemorySnapshotsPartition(parentTable.commits.toSeq.map(c => (c.id, c.id))))
  }

  override def createReaderFactory(): PartitionReaderFactory = {
    new InMemorySnapshotsReaderFactory()
  }
}

case class InMemorySnapshotsPartition(snapshots: Seq[(Long, Long)]) extends InputPartition

class InMemorySnapshotsReaderFactory extends PartitionReaderFactory {
  override def createReader(partition: InputPartition): PartitionReader[InternalRow] = {
    new InMemorySnapshotsReader(partition.asInstanceOf[InMemorySnapshotsPartition])
  }
}

class InMemorySnapshotsReader(partition: InMemorySnapshotsPartition)
    extends PartitionReader[InternalRow] {
  private var index = -1
  private val snapshots = partition.snapshots

  override def next(): Boolean = {
    index += 1
    index < snapshots.size
  }

  override def get(): InternalRow = {
    val (committedAt, snapshotId) = snapshots(index)
    InternalRow(committedAt, snapshotId)
  }

  override def close(): Unit = {}
}
