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

import org.apache.hudi.common.schema.{HoodieSchema, HoodieSchemaType}

import org.apache.spark.sql.{Column, DataFrame, SparkSession}
import org.apache.spark.sql.functions.{lit, struct}
import org.apache.spark.sql.types._
import org.junit.jupiter.api.Assertions.{assertEquals, assertTrue}

import java.io.File
import java.nio.file.{Files, Path}

object BlobTestHelpers {
  def blobMetadata: Metadata = {
    new MetadataBuilder()
      .putString(HoodieSchema.TYPE_METADATA_FIELD, HoodieSchemaType.BLOB.name())
      .build()
  }

  def inlineBlobStructCol(name: String, bytesCol: Column): Column = {
    struct(
      lit(HoodieSchema.Blob.INLINE).as(HoodieSchema.Blob.TYPE),
      bytesCol.cast("binary").as(HoodieSchema.Blob.INLINE_DATA_FIELD),
      lit(null).cast("struct<external_path:string,offset:bigint,length:bigint,managed:boolean>")
        .as(HoodieSchema.Blob.EXTERNAL_REFERENCE)
    ).as(name, blobMetadata)
  }

  def wholeFileBlobStructCol(name: String, filePathCol: Column): Column = {
    struct(
      lit(HoodieSchema.Blob.OUT_OF_LINE).as(HoodieSchema.Blob.TYPE),
      lit(null).cast("binary").as(HoodieSchema.Blob.INLINE_DATA_FIELD),
      struct(
        filePathCol.as(HoodieSchema.Blob.EXTERNAL_REFERENCE_PATH),
        lit(null).cast("bigint").as(HoodieSchema.Blob.EXTERNAL_REFERENCE_OFFSET),
        lit(null).cast("bigint").as(HoodieSchema.Blob.EXTERNAL_REFERENCE_LENGTH),
        lit(false).as(HoodieSchema.Blob.EXTERNAL_REFERENCE_IS_MANAGED)
      ).as(HoodieSchema.Blob.EXTERNAL_REFERENCE)
    ).as(name, blobMetadata)
  }

  def blobStructCol(name: String, filePathCol: Column, offsetCol: Column, lengthCol: Column): Column = {
    struct(
      lit(HoodieSchema.Blob.OUT_OF_LINE).as(HoodieSchema.Blob.TYPE),
      lit(null).cast("binary").as(HoodieSchema.Blob.INLINE_DATA_FIELD),
      struct(
        filePathCol.as(HoodieSchema.Blob.EXTERNAL_REFERENCE_PATH),
        offsetCol.as(HoodieSchema.Blob.EXTERNAL_REFERENCE_OFFSET),
        lengthCol.cast("bigint").as(HoodieSchema.Blob.EXTERNAL_REFERENCE_LENGTH),
        lit(false).as(HoodieSchema.Blob.EXTERNAL_REFERENCE_IS_MANAGED)
      ).as(HoodieSchema.Blob.EXTERNAL_REFERENCE)
    ).as(name, blobMetadata)
  }

  def createTestFile(tempDir: Path, name: String, size: Int): String = {
    val file = new File(tempDir.toString, name)
    val bytes = (0 until size).map(i => (i % 256).toByte).toArray
    Files.write(file.toPath, bytes)
    file.getAbsolutePath
  }

  /**
   * Assert that byte array contains expected pattern (i % 256) at given offset.
   *
   * @param data Array of bytes to verify
   * @param expectedOffset Starting offset for pattern (default 0)
   */
  def assertBytesContent(data: Array[Byte], expectedOffset: Int = 0): Unit = {
    for (i <- 0 until data.length) {
      assertEquals((expectedOffset + i) % 256, data(i) & 0xFF,
        s"Mismatch at byte $i (global offset ${expectedOffset + i})")
    }
  }

  /**
   * Execute code block with temporary Spark configuration.
   * Automatically restores previous values after execution.
   *
   * @param spark SparkSession instance
   * @param configs Configuration key-value pairs to set
   * @param fn Code block to execute with configs
   */
  def withSparkConfig[T](spark: SparkSession, configs: Map[String, String])(fn: => T): T = {
    val oldValues = configs.keys.map(k => (k, spark.conf.getOption(k))).toMap
    try {
      configs.foreach { case (k, v) => spark.conf.set(k, v) }
      fn
    } finally {
      oldValues.foreach { case (k, oldValue) =>
        oldValue match {
          case Some(v) => spark.conf.set(k, v)
          case None => spark.conf.unset(k)
        }
      }
    }
  }

  /**
   * Assert that DataFrame contains all specified columns.
   *
   * @param df DataFrame to check
   * @param columnNames Variable number of column names
   */
  def assertColumnsExist(df: DataFrame, columnNames: String*): Unit = {
    columnNames.foreach { colName =>
      assertTrue(df.columns.contains(colName),
        s"DataFrame missing expected column: $colName. Available: ${df.columns.mkString(", ")}")
    }
  }
}
