/*
 * 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.hudi.command.procedures

import org.apache.hudi.client.transaction.lock.audit.StorageLockProviderAuditService
import org.apache.hudi.util.JFunction

import org.apache.hadoop.fs.{Path, FileStatus}

import com.fasterxml.jackson.databind.ObjectMapper
import org.apache.spark.sql.Row
import org.apache.spark.sql.types.{DataTypes, Metadata, StructField, StructType}

import java.io.FileNotFoundException
import java.util.function.Supplier

import scala.collection.mutable
import scala.util.{Failure, Success, Try}

/**
 * Spark SQL procedure for validating audit lock files for Hudi tables.
 *
 * This procedure validates the integrity and consistency of audit lock files
 * generated by the storage lock audit service. It checks for various issues
 * including file format, lock state transitions, and orphaned locks.
 *
 * Usage:
 * {{{
 * CALL validate_audit_lock(table => 'my_table')
 * CALL validate_audit_lock(path => '/path/to/table')
 * }}}
 *
 * The procedure reads audit files from:
 * `{table_path}/.hoodie/.locks/audit/`
 *
 * @author Apache Hudi
 * @since 1.0.0
 */
class ValidateAuditLockProcedure extends BaseProcedure with ProcedureBuilder {
  private val PARAMETERS = Array[ProcedureParameter](
    ProcedureParameter.optional(0, "table", DataTypes.StringType),
    ProcedureParameter.optional(1, "path", DataTypes.StringType)
  )

  private val OUTPUT_TYPE = new StructType(Array[StructField](
    StructField("table", DataTypes.StringType, nullable = false, Metadata.empty),
    StructField("validation_result", DataTypes.StringType, nullable = false, Metadata.empty),
    StructField("transactions_validated", DataTypes.IntegerType, nullable = false, Metadata.empty),
    StructField("issues_found", DataTypes.IntegerType, nullable = false, Metadata.empty),
    StructField("details", DataTypes.StringType, nullable = false, Metadata.empty)
  ))

  private val OBJECT_MAPPER = new ObjectMapper()

  /**
   * Represents a transaction window with start time, end time, and metadata.
   */
  case class TransactionWindow(
    ownerId: String,
    transactionCreationTime: Long,
    lockAcquisitionTime: Long,
    lockReleaseTime: Option[Long],
    lockExpirationTime: Option[Long],
    filename: String
  ) {
    def effectiveEndTime: Long = lockReleaseTime.orElse(lockExpirationTime).getOrElse(lockAcquisitionTime)
  }

  /**
   * Returns the procedure parameters definition.
   *
   * @return Array of parameters: table (optional String) and path (optional String)
   */
  def parameters: Array[ProcedureParameter] = PARAMETERS

  /**
   * Returns the output schema for the procedure result.
   *
   * @return StructType containing table, validation_result, issues_found, and details columns
   */
  def outputType: StructType = OUTPUT_TYPE

  /**
   * Executes the audit lock validation procedure.
   *
   * @param args Procedure arguments containing table name or path
   * @return Sequence containing a single Row with validation results
   * @throws IllegalArgumentException if neither table nor path is provided, or both are provided
   */
  override def call(args: ProcedureArgs): Seq[Row] = {
    super.checkArgs(PARAMETERS, args)

    val tableName = getArgValueOrDefault(args, PARAMETERS(0))
    val tablePath = getArgValueOrDefault(args, PARAMETERS(1))

    // Get the base path using BaseProcedure helper (handles table/path validation)
    val basePath: String = getBasePath(tableName, tablePath)
    val metaClient = createMetaClient(jsc, basePath)

    // Use table name if provided, otherwise extract from path
    val displayName = tableName.map(_.asInstanceOf[String]).getOrElse(tablePath.get.asInstanceOf[String])

    try {
      val auditFolderPath = new Path(StorageLockProviderAuditService.getAuditFolderPath(basePath))
      val fs = metaClient.getFs

      // Get all audit files
      val allFilesResult = try {
        Some(fs.listStatus(auditFolderPath))
      } catch {
        case _: FileNotFoundException =>
          None
      }

      allFilesResult match {
        case None =>
          Seq(Row(displayName, "PASSED", 0, 0, "No audit folder found - nothing to validate"))
        case Some(allFiles) =>
          val auditFiles = allFiles.filter(pathInfo => pathInfo.isFile && pathInfo.getPath.getName.endsWith(".jsonl"))

          if (auditFiles.isEmpty) {
            Seq(Row(displayName, "PASSED", 0, 0, "No audit files found - nothing to validate"))
          } else {

            // Parse all audit files into transaction windows
            val parseResults = auditFiles.map(pathInfo => (pathInfo, parseAuditFile(pathInfo, fs)))
            val windows = parseResults.flatMap(_._2).toSeq
            val parseErrors = parseResults.filter(_._2.isEmpty).map(p =>
              s"[ERROR] Failed to parse audit file: ${p._1.getPath.getName}")

            if (windows.isEmpty) {
              Seq(Row(displayName, "FAILED", 0, auditFiles.size, "Failed to parse any audit files"))
            } else {

              // Validate transactions
              val validationResults = validateTransactionWindows(windows)

              // Add parse errors to validation results
              val allErrors = validationResults.errors ++ parseErrors
              val resultsWithParseErrors = validationResults.copy(errors = allErrors.toList)

              // Generate result
              val (result, issuesFound, details) = formatValidationResults(resultsWithParseErrors, auditFiles.size, windows.size, parseErrors.size)

              Seq(Row(displayName, result, windows.size, issuesFound, details))
            }
          }
      }
    } catch {
      case e: Exception =>
        val errorMessage = s"Validation failed: ${e.getMessage}"
        Seq(Row(displayName, "ERROR", 0, -1, errorMessage))
    }
  }

  /**
   * Parses an audit file and extracts transaction window information.
   */
  private def parseAuditFile(pathInfo: FileStatus, fs: org.apache.hadoop.fs.FileSystem): Option[TransactionWindow] = {
    val filename = pathInfo.getPath.getName

    Try {
      // Read and parse JSONL content
      val inputStream = fs.open(pathInfo.getPath)
      val jsonNodes = try {
        val content = scala.io.Source.fromInputStream(inputStream).mkString
        val lines = content.split('\n').filter(_.trim.nonEmpty)
        lines.map(line => OBJECT_MAPPER.readTree(line)).toSeq
      } finally {
        inputStream.close()
      }

      if (jsonNodes.isEmpty) {
        None
      } else {

        // Extract transaction metadata from first entry
        val firstNode = jsonNodes.head
        val ownerId = firstNode.get("ownerId").asText()
        val transactionCreationTime = firstNode.get("transactionStartTime").asLong()

        // Find first START timestamp
        val startNodes = jsonNodes.filter(node => node.get("state").asText() == "START")
        val lockAcquisitionTime = if (startNodes.nonEmpty) {
          startNodes.head.get("timestamp").asLong()
        } else {
          transactionCreationTime // fallback to transaction creation time
        }

        // Find last END timestamp
        val endNodes = jsonNodes.filter(node => node.get("state").asText() == "END")
        val lockReleaseTime = if (endNodes.nonEmpty) {
          Some(endNodes.last.get("timestamp").asLong())
        } else {
          None
        }

        // Find last expiration time as fallback
        val lockExpirationTime = if (jsonNodes.nonEmpty) {
          Some(jsonNodes.last.get("lockExpiration").asLong())
        } else {
          None
        }

        Some(TransactionWindow(
          ownerId = ownerId,
          transactionCreationTime = transactionCreationTime,
          lockAcquisitionTime = lockAcquisitionTime,
          lockReleaseTime = lockReleaseTime,
          lockExpirationTime = lockExpirationTime,
          filename = filename
        ))
      }
    } match {
      case Success(window) => window
      case Failure(_) => None // Skip corrupted files
    }
  }

  /**
   * Validates transaction windows for overlaps and proper closure.
   */
  private def validateTransactionWindows(windows: Seq[TransactionWindow]): ValidationResults = {
    val errors = mutable.ListBuffer[String]()
    val warnings = mutable.ListBuffer[String]()

    // Check for transactions without proper END
    windows.foreach { window =>
      if (window.lockReleaseTime.isEmpty) {
        warnings += s"[WARNING] ${window.filename} => transaction did not end gracefully. This could be due to driver OOM or non-graceful shutdown."
      }
    }

    // Sort windows by start time for overlap detection
    val sortedWindows = windows.sortBy(_.lockAcquisitionTime)

    // Check for overlaps
    for (i <- sortedWindows.indices) {
      val currentWindow = sortedWindows(i)
      val currentEnd = currentWindow.effectiveEndTime

      // Check all subsequent windows for overlaps
      for (j <- (i + 1) until sortedWindows.length) {
        val otherWindow = sortedWindows(j)
        val otherStart = otherWindow.lockAcquisitionTime

        // Check if windows overlap
        if (otherStart < currentEnd) {
          errors += s"[ERROR] ${currentWindow.filename} => overlaps with ${otherWindow.filename}"
        }
      }
    }

    ValidationResults(errors.toList, warnings.toList)
  }

  /**
   * Formats validation results into the output format.
   */
  private def formatValidationResults(results: ValidationResults, totalFiles: Int, parsedFiles: Int, failedFiles: Int): (String, Int, String) = {
    val totalIssues = results.errors.size + results.warnings.size

    if (totalIssues == 0) {
      ("PASSED", 0, s"All audit lock transactions validated successfully. Audit Files: $totalFiles total, $parsedFiles parsed successfully, $failedFiles failed to parse")
    } else {
      val resultType = if (results.errors.nonEmpty) "FAILED" else "WARNING"
      val fileInfo = s"Audit Files: $totalFiles total, $parsedFiles parsed successfully, $failedFiles failed to parse. "
      val details = fileInfo + (results.errors ++ results.warnings).mkString(", ")
      (resultType, totalIssues, details)
    }
  }

  /**
   * Case class to hold validation results.
   */
  private case class ValidationResults(errors: List[String], warnings: List[String])

  /**
   * Builds a new instance of the ValidateAuditLockProcedure.
   *
   * @return New ValidateAuditLockProcedure instance
   */
  override def build: Procedure = new ValidateAuditLockProcedure()
}


/**
 * Companion object for ValidateAuditLockProcedure containing constants and factory methods.
 */
object ValidateAuditLockProcedure {
  /** The name used to register and invoke this procedure */
  val NAME = "validate_audit_lock"

  /**
   * Factory method to create procedure builder instances.
   *
   * @return Supplier that creates new ValidateAuditLockProcedure instances
   */
  def builder: Supplier[ProcedureBuilder] = JFunction.toJavaSupplier(() => new ValidateAuditLockProcedure())
}
