/*
 * 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.auron.arrowio

import java.lang.Thread.UncaughtExceptionHandler
import java.security.PrivilegedExceptionAction
import java.util.concurrent.ArrayBlockingQueue
import java.util.concurrent.BlockingQueue

import org.apache.arrow.c.ArrowArray
import org.apache.arrow.c.ArrowSchema
import org.apache.arrow.c.Data
import org.apache.arrow.vector.VectorSchemaRoot
import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider
import org.apache.spark.TaskContext
import org.apache.spark.internal.Logging
import org.apache.spark.sql.auron.NativeHelper
import org.apache.spark.sql.auron.util.Using
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.execution.auron.arrowio.util.ArrowUtils
import org.apache.spark.sql.execution.auron.arrowio.util.ArrowUtils.CHILD_ALLOCATOR
import org.apache.spark.sql.execution.auron.arrowio.util.ArrowUtils.ROOT_ALLOCATOR
import org.apache.spark.sql.execution.auron.arrowio.util.ArrowWriter
import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.vectorized.ColumnarBatch

import org.apache.auron.arrowio.AuronArrowFFIExporter
import org.apache.auron.configuration.AuronConfiguration
import org.apache.auron.jni.AuronAdaptor

class ArrowFFIExporter(inputIter: Iterator[Any], schema: StructType)
    extends AuronArrowFFIExporter
    with Logging {
  private val sparkAuronConfig: AuronConfiguration =
    AuronAdaptor.getInstance.getAuronConfiguration
  private val maxBatchNumRows = sparkAuronConfig.getInteger(AuronConfiguration.BATCH_SIZE)
  private val maxBatchMemorySize =
    sparkAuronConfig.getInteger(AuronConfiguration.SUGGESTED_BATCH_MEM_SIZE)

  private val arrowSchema = ArrowUtils.toArrowSchema(schema)
  private val emptyDictionaryProvider = new MapDictionaryProvider()
  private val nativeCurrentUser = NativeHelper.currentUser

  private trait QueueState
  private case object NextBatch extends QueueState
  private case class Finished(t: Option[Throwable]) extends QueueState

  private val tc = TaskContext.get()
  // Build a meaningful identifier from TaskContext info
  private val exporterId = if (tc != null) {
    s"stage-${tc.stageId()}-part-${tc.partitionId()}-tid-${tc.taskAttemptId()}-${System.identityHashCode(this)}"
  } else {
    s"no-context-${System.identityHashCode(this)}"
  }
  private val closed = new java.util.concurrent.atomic.AtomicBoolean(false)
  private val outputQueue: BlockingQueue[QueueState] = new ArrayBlockingQueue[QueueState](16)
  private val processingQueue: BlockingQueue[Unit] = new ArrayBlockingQueue[Unit](16)
  private var currentRoot: VectorSchemaRoot = _
  private val rowIter = new InputToRowIter(inputIter)
  private val outputThread = startOutputThread()

  def exportSchema(exportArrowSchemaPtr: Long): Unit = {
    Using.resource(ArrowSchema.wrap(exportArrowSchemaPtr)) { exportSchema =>
      Data.exportSchema(ROOT_ALLOCATOR, arrowSchema, emptyDictionaryProvider, exportSchema)
    }
  }

  override def exportNextBatch(exportArrowArrayPtr: Long): Boolean = {
    if (!hasNext) {
      return false
    }

    // export using root allocator
    Using.resource(ArrowArray.wrap(exportArrowArrayPtr)) { exportArray =>
      Data.exportVectorSchemaRoot(
        ROOT_ALLOCATOR,
        currentRoot,
        emptyDictionaryProvider,
        exportArray)
    }

    // to continue processing next batch
    processingQueue.put(())
    true
  }

  private def hasNext: Boolean = {
    if (tc != null && (tc.isCompleted() || tc.isInterrupted())) {
      return false
    }
    outputQueue.take() match {
      case NextBatch => true
      case Finished(None) => false
      case Finished(Some(e)) => throw e
    }
  }

  private def startOutputThread(): Thread = {
    val thread = new Thread(new Runnable {
      override def run(): Unit = {
        if (tc != null) {
          TaskContext.setTaskContext(tc)
        }

        nativeCurrentUser.doAs(new PrivilegedExceptionAction[Unit] {
          override def run(): Unit = {
            try {
              while (tc == null || (!tc.isCompleted() && !tc.isInterrupted())) {
                if (!rowIter.hasNext) {
                  outputQueue.put(Finished(None))
                  return
                }

                Using.resource(CHILD_ALLOCATOR("ArrowFFIExporter")) { allocator =>
                  Using.resource(VectorSchemaRoot.create(arrowSchema, allocator)) { root =>
                    val arrowWriter = ArrowWriter.create(root)
                    while (rowIter.hasNext
                      && allocator.getAllocatedMemory < maxBatchMemorySize
                      && arrowWriter.currentCount < maxBatchNumRows) {
                      arrowWriter.write(rowIter.next())
                    }
                    arrowWriter.finish()

                    // export root
                    currentRoot = root
                    outputQueue.put(NextBatch)

                    // wait for processing next batch
                    processingQueue.take()
                  }
                }
              }
              outputQueue.put(Finished(None))
            } catch {
              case _: InterruptedException =>
                // Thread was interrupted during close(), this is expected - just exit gracefully
                logDebug(s"ArrowFFIExporter-$exporterId: outputThread interrupted, exiting")
                outputQueue.clear()
                outputQueue.put(Finished(None))
            } finally {
              rowIter.close()
            }
          }
        })
      }
    })

    if (tc != null) {
      tc.addTaskCompletionListener[Unit]((_: TaskContext) => close())
      tc.addTaskFailureListener((_, _) => close())
    }

    thread.setName(s"ArrowFFIExporter-$exporterId")
    thread.setDaemon(true)
    thread.setUncaughtExceptionHandler(new UncaughtExceptionHandler {
      override def uncaughtException(t: Thread, e: Throwable): Unit = {
        outputQueue.clear()
        outputQueue.put(Finished(Some(e)))
      }
    })
    thread.start()
    thread
  }

  override def close(): Unit = {
    // Ensure close() is idempotent - only execute once
    if (!closed.compareAndSet(false, true)) {
      logDebug(s"ArrowFFIExporter-$exporterId: close() already called, skipping")
      return
    }

    if (outputThread.isAlive) {
      logDebug(s"ArrowFFIExporter-$exporterId: interrupting outputThread")
      outputThread.interrupt()
      // Wait for the thread to terminate to ensure resources are properly released
      try {
        outputThread.join(5000) // Wait up to 5 seconds
        if (outputThread.isAlive) {
          logWarning(
            s"ArrowFFIExporter-$exporterId: outputThread did not terminate within 5 seconds")
        }
      } catch {
        case _: InterruptedException =>
          // Ignore - we don't need to propagate this to caller
          logDebug(s"ArrowFFIExporter-$exporterId: interrupted while waiting for outputThread")
      }
      logDebug(s"ArrowFFIExporter-$exporterId: close() completed")
    }
  }

  private class InputToRowIter(inputIter: Iterator[Any]) extends Iterator[InternalRow] {
    private var currentBatch: ColumnarBatch = _
    private var currentBatchRowId = 0
    private var pendingRow: InternalRow = _

    override def hasNext: Boolean = {
      if (pendingRow != null) {
        return true
      }

      closeFinishedBatch()
      if (currentBatch != null) {
        return true
      }

      while (inputIter.hasNext) {
        inputIter.next() match {
          case row: InternalRow =>
            pendingRow = row
            return true
          case batch: ColumnarBatch if batch.numRows() > 0 =>
            currentBatch = batch
            currentBatchRowId = 0
            return true
          case batch: ColumnarBatch =>
            batch.close()
          case null =>
            throw new IllegalStateException(
              "ArrowFFIExporter expects InternalRow or ColumnarBatch input, but got null")
          case other =>
            throw new IllegalStateException(
              s"ArrowFFIExporter expects InternalRow or ColumnarBatch input, " +
                s"but got ${other.getClass.getName}")
        }
      }

      false
    }

    override def next(): InternalRow = {
      if (!hasNext) {
        throw new NoSuchElementException("no more rows")
      }

      if (pendingRow != null) {
        val row = pendingRow
        pendingRow = null
        row
      } else {
        val row = currentBatch.getRow(currentBatchRowId)
        currentBatchRowId += 1
        row
      }
    }

    def close(): Unit = {
      if (currentBatch != null) {
        currentBatch.close()
        currentBatch = null
      }
      currentBatchRowId = 0
      pendingRow = null
    }

    private def closeFinishedBatch(): Unit = {
      if (currentBatch != null && currentBatchRowId >= currentBatch.numRows()) {
        currentBatch.close()
        currentBatch = null
        currentBatchRowId = 0
      }
    }
  }
}
