/*
 * 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.texera.amber.core.storage.result.iceberg

import org.apache.texera.amber.core.storage.IcebergCatalogInstance
import org.apache.texera.amber.core.storage.model.{BufferedItemWriter, VirtualDocument}
import org.apache.texera.amber.core.storage.util.StorageUtil.{withLock, withReadLock, withWriteLock}
import org.apache.texera.amber.util.IcebergUtil
import org.apache.commons.io.IOUtils
import org.apache.iceberg.catalog.{Catalog, TableIdentifier}
import org.apache.iceberg.data.Record
import org.apache.iceberg.exceptions.NoSuchTableException
import org.apache.iceberg.types.{Conversions, Types}
import org.apache.iceberg.{FileScanTask, Table}

import java.io.{ByteArrayInputStream, InputStream, PipedInputStream, PipedOutputStream}
import java.net.URI
import java.nio.ByteBuffer
import java.time.format.DateTimeFormatter
import java.time.{Instant, LocalDate, ZoneOffset}
import java.util.concurrent.locks.{ReentrantLock, ReentrantReadWriteLock}
import java.util.zip.{ZipEntry, ZipOutputStream}
import scala.collection.mutable
import scala.jdk.CollectionConverters._
import scala.util.Using

object Constants {
  // Default buffer size for streaming Parquet files to outside
  val DEFAULT_BUFFER_SIZE: Int = 8192
}

/**
  * IcebergDocument is used to read and write a set of T as an Iceberg table.
  * It provides iterator-based read methods and supports multiple writers to write to the same table.
  *
  * The table must exist when constructing the document
  *
  * @param tableNamespace namespace of the table.
  * @param tableName      name of the table.
  * @param tableSchema    schema of the table.
  * @param serde          function to serialize T into an Iceberg Record.
  * @param deserde        function to deserialize an Iceberg Record into T.
  * @tparam T type of the data items stored in the Iceberg table.
  */
private[storage] class IcebergDocument[T >: Null <: AnyRef](
    val tableNamespace: String,
    val tableName: String,
    val tableSchema: org.apache.iceberg.Schema,
    val serde: (org.apache.iceberg.Schema, T) => Record,
    val deserde: (org.apache.iceberg.Schema, Record) => T
) extends VirtualDocument[T]
    with OnIceberg {

  private val lock = new ReentrantReadWriteLock()

  @transient lazy val catalog: Catalog = IcebergCatalogInstance.getInstance()

  /**
    * Returns the URI of the table location.
    *
    * @throws NoSuchTableException if the table does not exist.
    */
  override def getURI: URI = {
    val table = IcebergUtil
      .loadTableMetadata(catalog, tableNamespace, tableName)
      .getOrElse(
        throw new NoSuchTableException(f"table ${tableNamespace}.${tableName} doesn't exist")
      )
    URI.create(table.location())
  }

  /**
    * Deletes the table and clears its contents.
    */
  override def clear(): Unit =
    withWriteLock(lock) {
      val identifier = TableIdentifier.of(tableNamespace, tableName)
      if (catalog.tableExists(identifier)) {
        catalog.dropTable(identifier)
      }
    }

  /**
    * Get an iterator for reading records from the table.
    */
  override def get(): Iterator[T] = getUsingFileSequenceOrder(0, None)

  /**
    * Get records within a specified range [from, until).
    */
  override def getRange(from: Int, until: Int, columns: Option[Seq[String]] = None): Iterator[T] = {
    getUsingFileSequenceOrder(from, Some(until), columns)
  }

  /**
    * Get records starting after a specified offset.
    */
  override def getAfter(offset: Int): Iterator[T] = {
    getUsingFileSequenceOrder(offset, None)
  }

  /**
    * Get the total count of records in the table.
    */
  override def getCount: Long = {
    val table = IcebergUtil
      .loadTableMetadata(catalog, tableNamespace, tableName)
      .getOrElse(
        return 0
      )
    table.newScan().planFiles().iterator().asScala.map(f => f.file().recordCount()).sum
  }

  /**
    * Creates a BufferedItemWriter for writing data to the table.
    *
    * @param writerIdentifier The writer's ID. It should be unique within the same table, as each writer will use it as
    *                         the prefix of the files they append
    */
  override def writer(writerIdentifier: String): BufferedItemWriter[T] = {
    new IcebergTableWriter[T](
      writerIdentifier,
      catalog,
      tableNamespace,
      tableName,
      tableSchema,
      serde
    )
  }

  /**
    * Util iterator to get T in certain range
    *
    * @param from  start from which record inclusively, if 0 means start from the first
    * @param until end at which record exclusively, if None means read to the table's EOF
    * @param columns columns to be projected
    */
  private def getUsingFileSequenceOrder(
      from: Int,
      until: Option[Int],
      columns: Option[Seq[String]] = None
  ): Iterator[T] =
    withReadLock(lock) {
      new Iterator[T] {
        private val iteLock = new ReentrantLock()
        // Load the table instance, initially the table instance may not exist
        private var table: Option[Table] = loadTableMetadata()

        // Last seen snapshot id(logically it's like a version number). While reading, new snapshots may be created
        private var lastSnapshotId: Option[Long] = None

        // Counter for how many records have been skipped
        private var numOfSkippedRecords = 0

        // Counter for how many records have been returned
        private var numOfReturnedRecords = 0

        // Total number of records to return
        private val totalRecordsToReturn = until.map(_ - from).getOrElse(Int.MaxValue)

        // Iterator for usable file scan tasks
        private var usableFileIterator: Iterator[FileScanTask] = seekToUsableFile()

        // Current record iterator for the active file
        private var currentRecordIterator: Iterator[Record] = Iterator.empty

        // Util function to load the table's metadata
        private def loadTableMetadata(): Option[Table] = {
          IcebergUtil.loadTableMetadata(
            catalog,
            tableNamespace,
            tableName
          )
        }

        private def seekToUsableFile(): Iterator[FileScanTask] =
          withLock(iteLock) {
            if (numOfSkippedRecords > from) {
              throw new RuntimeException("seek operation should not be called")
            }

            // refresh the table's snapshots
            if (table.isEmpty) {
              table = loadTableMetadata()
            }
            table.foreach(_.refresh())

            // Retrieve and sort the file scan tasks by file sequence number
            val fileScanTasksIterator: Iterator[FileScanTask] = table match {
              case Some(t) =>
                val currentSnapshotId = Option(t.currentSnapshot()).map(_.snapshotId())
                val fileScanTasks = (lastSnapshotId, currentSnapshotId) match {
                  // Read from the start
                  case (None, Some(_)) =>
                    val tasks = t.newScan().planFiles().iterator().asScala
                    lastSnapshotId = currentSnapshotId
                    tasks

                  // Read incrementally from the last snapshot
                  case (Some(lastId), Some(currId)) if lastId != currId =>
                    val tasks = t
                      .newIncrementalAppendScan()
                      .fromSnapshotExclusive(lastId)
                      .toSnapshot(currId)
                      .planFiles()
                      .iterator()
                      .asScala
                    lastSnapshotId = currentSnapshotId
                    tasks

                  // No new data
                  case (Some(lastId), Some(currId)) if lastId == currId =>
                    Iterator.empty

                  // Default: No data yet
                  case _ =>
                    Iterator.empty
                }
                fileScanTasks.toSeq.sortBy(_.file().fileSequenceNumber()).iterator

              case None =>
                Iterator.empty
            }

            // Iterate through sorted FileScanTasks and update numOfSkippedRecords
            val usableTasks = fileScanTasksIterator.dropWhile { task =>
              val recordCount = task.file().recordCount()
              if (numOfSkippedRecords + recordCount <= from) {
                numOfSkippedRecords += recordCount.toInt
                true
              } else {
                false
              }
            }

            usableTasks
          }

        override def hasNext: Boolean = {
          if (numOfReturnedRecords >= totalRecordsToReturn) {
            return false
          }

          // If the current file still has records, return immediately without touching the catalog.
          // Only when the current file is exhausted do we check for more files and possibly refresh.
          if (currentRecordIterator.hasNext) {
            return true
          }

          if (!usableFileIterator.hasNext) {
            usableFileIterator = seekToUsableFile()
          }

          while (!currentRecordIterator.hasNext && usableFileIterator.hasNext) {
            val nextFile = usableFileIterator.next()
            val schemaToUse = columns match {
              case Some(cols) => tableSchema.select(cols.asJava)
              case None       => tableSchema
            }
            currentRecordIterator = IcebergUtil.readDataFileAsIterator(
              nextFile.file(),
              schemaToUse,
              table.get
            )

            // Skip records within the file if necessary
            val recordsToSkipInFile = from - numOfSkippedRecords
            if (recordsToSkipInFile > 0) {
              currentRecordIterator = currentRecordIterator.drop(recordsToSkipInFile)
              numOfSkippedRecords += recordsToSkipInFile
            }
          }

          currentRecordIterator.hasNext
        }

        override def next(): T = {
          if (!hasNext) throw new NoSuchElementException("No more records available")

          val record = currentRecordIterator.next()
          numOfReturnedRecords += 1
          val schemaToUse = columns match {
            case Some(cols) => tableSchema.select(cols.asJava)
            case None       => tableSchema
          }
          deserde(schemaToUse, record)
        }
      }
    }

  /**
    * Provides methods for extracting metadata statistics for the results
    *
    * **Statistics Computed:**
    * - **Numeric fields (Int, Long, Double)**: Computes `min` and `max`.
    * - **Date fields (Timestamp)**: Computes `min` and `max` (converted to `LocalDate`).
    * - **All fields**: Computes `not_null_count` (number of non-null values).
    *
    * @return A map where each field name is mapped to a nested map containing its statistics.
    * @throws NoSuchTableException if the table does not exist in the catalog.
    */
  override def getTableStatistics: Map[String, Map[String, Any]] = {
    val table = IcebergUtil
      .loadTableMetadata(catalog, tableNamespace, tableName)
      .getOrElse(
        throw new NoSuchTableException(f"table ${tableNamespace}.${tableName} doesn't exist")
      )

    val schema = table.schema()

    // Extract field names, IDs, and types from the schema
    val fieldTypes =
      schema.columns().asScala.map(col => col.name() -> (col.fieldId(), col.`type`())).toMap
    val fieldStats = mutable.Map[String, mutable.Map[String, Any]]()
    val dateFormatter: DateTimeFormatter = DateTimeFormatter.ISO_LOCAL_DATE

    // Initialize statistics for each field
    fieldTypes.foreach {
      case (field, (_, fieldType)) =>
        val initialStats = mutable.Map[String, Any](
          "not_null_count" -> 0L
        )
        if (
          fieldType == Types.IntegerType.get() || fieldType == Types.LongType
            .get() || fieldType == Types.DoubleType.get()
        ) {
          initialStats("min") = Double.MaxValue
          initialStats("max") = Double.MinValue
        } else if (
          fieldType == Types.TimestampType.withoutZone() || fieldType == Types.TimestampType
            .withZone()
        ) {
          initialStats("min") = LocalDate.MAX.format(dateFormatter)
          initialStats("max") = LocalDate.MIN.format(dateFormatter)
        }

        fieldStats(field) = initialStats
    }

    // Scan table files and aggregate statistics
    table.newScan().includeColumnStats().planFiles().iterator().asScala.foreach { file =>
      val fileStats = file.file()
      // Extract column-level statistics
      val lowerBounds =
        Option(fileStats.lowerBounds()).getOrElse(Map.empty[Integer, ByteBuffer].asJava)
      val upperBounds =
        Option(fileStats.upperBounds()).getOrElse(Map.empty[Integer, ByteBuffer].asJava)
      val nullCounts =
        Option(fileStats.nullValueCounts()).getOrElse(Map.empty[Integer, java.lang.Long].asJava)
      val nanCounts =
        Option(fileStats.nanValueCounts()).getOrElse(Map.empty[Integer, java.lang.Long].asJava)

      fieldTypes.foreach {
        case (field, (fieldId, fieldType)) =>
          val lowerBound = Option(lowerBounds.get(fieldId))
          val upperBound = Option(upperBounds.get(fieldId))
          val nullCount: Long = Option(nullCounts.get(fieldId)).map(_.toLong).getOrElse(0L)
          val nanCount: Long = Option(nanCounts.get(fieldId)).map(_.toLong).getOrElse(0L)
          val fieldStat = fieldStats(field)

          // Process min/max values for numerical types
          if (
            fieldType == Types.IntegerType.get() || fieldType == Types.LongType
              .get() || fieldType == Types.DoubleType.get()
          ) {
            lowerBound.foreach { buffer =>
              val minValue =
                Conversions.fromByteBuffer(fieldType, buffer).asInstanceOf[Number].doubleValue()
              fieldStat("min") = Math.min(fieldStat("min").asInstanceOf[Double], minValue)
            }

            upperBound.foreach { buffer =>
              val maxValue =
                Conversions.fromByteBuffer(fieldType, buffer).asInstanceOf[Number].doubleValue()
              fieldStat("max") = Math.max(fieldStat("max").asInstanceOf[Double], maxValue)
            }
          }
          // Process min/max values for timestamp types
          else if (
            fieldType == Types.TimestampType.withoutZone() || fieldType == Types.TimestampType
              .withZone()
          ) {
            lowerBound.foreach { buffer =>
              val epochMicros = Conversions
                .fromByteBuffer(Types.TimestampType.withoutZone(), buffer)
                .asInstanceOf[Long]
              val dateValue =
                Instant.ofEpochMilli(epochMicros / 1000).atZone(ZoneOffset.UTC).toLocalDate
              fieldStat("min") =
                if (
                  dateValue
                    .isBefore(LocalDate.parse(fieldStat("min").asInstanceOf[String], dateFormatter))
                )
                  dateValue.format(dateFormatter)
                else
                  fieldStat("min")
            }

            upperBound.foreach { buffer =>
              val epochMicros = Conversions
                .fromByteBuffer(Types.TimestampType.withoutZone(), buffer)
                .asInstanceOf[Long]
              val dateValue =
                Instant.ofEpochMilli(epochMicros / 1000).atZone(ZoneOffset.UTC).toLocalDate
              fieldStat("max") =
                if (
                  dateValue
                    .isAfter(LocalDate.parse(fieldStat("max").asInstanceOf[String], dateFormatter))
                )
                  dateValue.format(dateFormatter)
                else
                  fieldStat("max")
            }
          }
          // Update non-null count
          fieldStat("not_null_count") = fieldStat("not_null_count").asInstanceOf[Long] +
            (fileStats.recordCount().toLong - nullCount - nanCount)
      }
    }
    fieldStats.map {
      case (field, stats) =>
        field -> stats.toMap
    }.toMap
  }

  /**
    * Computes the total size of all data files in the Iceberg table.
    *
    * @return The total size of all data files in bytes.
    * @throws NoSuchTableException if the table does not exist.
    */
  override def getTotalFileSize: Long = {
    val table = IcebergUtil
      .loadTableMetadata(catalog, tableNamespace, tableName)
      .getOrElse(
        throw new NoSuchTableException(f"table ${tableNamespace}.${tableName} doesn't exist")
      )
    var filesSize: Long = 0L

    try {
      filesSize = table.currentSnapshot().summary().getOrDefault("total-files-size", "0").toLong
    } catch {
      case e: Exception => println(s"Failed to get total-files-size: ${e.getMessage}")
    }
    filesSize
  }

  private def getParquetFileStream(fileTask: FileScanTask): InputStream = {
    val file = fileTask.file()
    val table = catalog.loadTable(TableIdentifier.of(tableNamespace, tableName))
    table.io().newInputFile(file.location()).newStream()
  }

  /**
    * Converts Iceberg Parquet files to ZIP and return as a piped stream.
    * Each file in the table is added as a separate entry in the ZIP archive.
    *
    * @return An InputStream that streams the contents of the Iceberg table as a ZIP archive.
    */
  override def asInputStream(): InputStream = {
    val fileScanTasks: Seq[FileScanTask] = {
      val table = this.catalog.loadTable(TableIdentifier.of(this.tableNamespace, this.tableName))
      table.refresh()
      table.newScan().planFiles().iterator().asScala.toSeq
    }

    if (fileScanTasks.isEmpty) {
      return new ByteArrayInputStream(Array.emptyByteArray)
    }

    val pipeIn = new PipedInputStream(Constants.DEFAULT_BUFFER_SIZE)
    val pipeOut = new PipedOutputStream(pipeIn)

    // Start processing in a separate thread to avoid blocking
    val processingThread = new Thread(() => {
      try {
        Using.resource(new ZipOutputStream(pipeOut)) { zipOut =>
          for (i <- fileScanTasks.indices) {
            val fileTask = fileScanTasks(i)
            val entryName = s"part-${String.format("%05d", i)}.parquet"
            zipOut.putNextEntry(new ZipEntry(entryName))
            Using.resource(getParquetFileStream(fileTask)) { fileInputStream =>
              IOUtils.copy(fileInputStream, zipOut)
            }
            zipOut.closeEntry()
          }
        }
      } catch {
        case e: Exception => throw e
      } finally {
        pipeOut.close()
      }
    })

    // Set the thread as a daemon so
    // 1- It will terminate if the request stop
    // 2- It handle lifecycle of cleaning up the pipe resources
    processingThread.setDaemon(true)
    processingThread.start()

    pipeIn
  }
}
