/*
 * 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.io.storage;

import org.apache.hudi.client.SparkTaskContextSupplier;
import org.apache.hudi.common.config.HoodieStorageConfig;
import org.apache.hudi.common.model.HoodieKey;
import org.apache.hudi.common.model.HoodieRecord.HoodieRecordType;
import org.apache.hudi.common.schema.HoodieSchema;
import org.apache.hudi.common.testutils.DisableDictionaryInjector;
import org.apache.hudi.config.HoodieWriteConfig;
import org.apache.hudi.io.HoodieParquetConfigInjector;
import org.apache.hudi.storage.StoragePath;
import org.apache.hudi.table.HoodieSparkTable;
import org.apache.hudi.table.HoodieTable;
import org.apache.hudi.testutils.HoodieClientTestBase;
import org.apache.hudi.testutils.SparkDatasetTestUtils;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.parquet.column.Encoding;
import org.apache.parquet.hadoop.ParquetFileReader;
import org.apache.parquet.hadoop.metadata.BlockMetaData;
import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData;
import org.apache.parquet.hadoop.metadata.ParquetMetadata;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.catalyst.InternalRow;
import org.apache.spark.sql.catalyst.expressions.GenericInternalRow;
import org.junit.jupiter.api.Test;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

/**
 * Tests for {@link HoodieParquetConfigInjector} functionality in {@link HoodieSparkFileWriterFactory}.
 */
public class TestHoodieParquetConfigInjector extends HoodieClientTestBase {

  /**
   * Get HoodieSchema that matches SparkDatasetTestUtils.STRUCT_TYPE.
   * This schema includes metadata fields and matches the structure of rows generated by getRandomRows().
   */
  private HoodieSchema getStructTypeCompatibleSchema() {
    // Create Avro schema that matches SparkDatasetTestUtils.STRUCT_TYPE
    // STRUCT_TYPE has: commit_time, commit_seqno, record_key, partition_path, filename,
    //                  RECORD_KEY_FIELD_NAME, PARTITION_PATH_FIELD_NAME, randomInt, randomLong
    String avroSchema = "{"
        + "\"type\":\"record\","
        + "\"name\":\"test_schema\","
        + "\"namespace\":\"test.namespace\","
        + "\"fields\":["
        + "{\"name\":\"_hoodie_commit_time\",\"type\":\"string\"},"
        + "{\"name\":\"_hoodie_commit_seqno\",\"type\":\"string\"},"
        + "{\"name\":\"_hoodie_record_key\",\"type\":\"string\"},"
        + "{\"name\":\"_hoodie_partition_path\",\"type\":\"string\"},"
        + "{\"name\":\"_hoodie_file_name\",\"type\":\"string\"},"
        + "{\"name\":\"key\",\"type\":\"string\"},"
        + "{\"name\":\"partition\",\"type\":\"string\"},"
        + "{\"name\":\"randomInt\",\"type\":\"int\"},"
        + "{\"name\":\"randomLong\",\"type\":\"long\"}"
        + "]}";
    return HoodieSchema.parse(avroSchema);
  }

  /**
   * Helper method to convert InternalRow (potentially UnsafeRow) to mutable GenericInternalRow.
   * This is needed because writeRowWithMetadata() calls row.update() which doesn't work on UnsafeRow.
   */
  private List<InternalRow> toMutableRows(List<InternalRow> internalRows) {
    List<InternalRow> mutableRows = new ArrayList<>();
    for (InternalRow row : internalRows) {
      Object[] values = new Object[row.numFields()];
      for (int i = 0; i < row.numFields(); i++) {
        if (row.isNullAt(i)) {
          values[i] = null;
        } else {
          values[i] = row.get(i, SparkDatasetTestUtils.STRUCT_TYPE.fields()[i].dataType());
        }
      }
      mutableRows.add(new GenericInternalRow(values));
    }
    return mutableRows;
  }

  @Test
  public void testDisableDictionaryEncodingViaInjector() throws Exception {
    final String instantTime = "100";
    final StoragePath parquetPath = new StoragePath(
        basePath + "/partition/path/test_dictionary_" + instantTime + ".parquet");

    HoodieSchema schema = getStructTypeCompatibleSchema();

    // Create config with the custom injector
    final HoodieWriteConfig cfg = getConfigBuilder()
        .withStorageConfig(HoodieStorageConfig.newBuilder()
            .withParquetConfigInjectorClass(DisableDictionaryInjector.class.getName())
            .parquetDictionaryEnabled(true) // Start with dictionary enabled
            .build())
        .build();

    HoodieTable table = HoodieSparkTable.create(cfg, context, metaClient);
    SparkTaskContextSupplier supplier = new SparkTaskContextSupplier();

    // Create writer and write some data
    HoodieFileWriter writer = HoodieFileWriterFactory.getFileWriter(
        instantTime, parquetPath, table.getStorage(), cfg.getStorageConfig(),
        schema, supplier, HoodieRecordType.SPARK);

    assertTrue(writer instanceof HoodieSparkParquetWriter);

    // Generate test data using SparkDatasetTestUtils and convert to InternalRows
    Dataset<Row> rowDataset = SparkDatasetTestUtils.getRandomRows(sqlContext, 100, "partition/path", false);
    List<InternalRow> internalRows = toMutableRows(SparkDatasetTestUtils.toInternalRows(rowDataset, SparkDatasetTestUtils.ENCODER));
    List<Row> rows = rowDataset.collectAsList();

    // Write some test records
    HoodieSparkParquetWriter sparkWriter = (HoodieSparkParquetWriter) writer;
    for (int i = 0; i < internalRows.size(); i++) {
      InternalRow row = internalRows.get(i);
      String recordKey = rows.get(i).getString(2); // record_key is at index 2 (after commit_time and seq_no)
      String partition = rows.get(i).getString(3);  // partition is at index 3
      sparkWriter.writeRowWithMetadata(new HoodieKey(recordKey, partition), row);
    }
    writer.close();

    // Verify the parquet file was created
    assertTrue(table.getStorage().exists(parquetPath));

    // Read parquet metadata and verify dictionary encoding is disabled
    Configuration hadoopConf = new Configuration();
    Path hadoopPath = new Path(parquetPath.toUri());
    ParquetFileReader reader = ParquetFileReader.open(hadoopConf, hadoopPath);
    ParquetMetadata metadata = reader.getFooter();
    reader.close();

    assertNotNull(metadata);

    // Verify that dictionary encoding is NOT used for any column
    // When dictionary encoding is disabled, columns should use PLAIN or other encodings but not RLE_DICTIONARY
    for (BlockMetaData block : metadata.getBlocks()) {
      for (ColumnChunkMetaData column : block.getColumns()) {
        // Check all encodings used for this column - should not include RLE_DICTIONARY or PLAIN_DICTIONARY
        for (Encoding encoding : column.getEncodings()) {
          assertFalse(encoding == Encoding.RLE_DICTIONARY || encoding == Encoding.PLAIN_DICTIONARY,
              "Column " + column.getPath() + " should not use dictionary encoding, but found: " + encoding);
        }
      }
    }
  }

  @Test
  public void testInvalidInjectorClassThrowsException() throws IOException {
    final String instantTime = "102";
    final StoragePath parquetPath = new StoragePath(
        basePath + "/partition/path/test_invalid_" + instantTime + ".parquet");

    HoodieSchema schema = getStructTypeCompatibleSchema();

    // Create config with an invalid/non-existent injector class
    final HoodieWriteConfig cfg = getConfigBuilder()
        .withStorageConfig(HoodieStorageConfig.newBuilder()
            .withParquetConfigInjectorClass("org.apache.hudi.NonExistentInjector")
            .build())
        .build();

    HoodieTable table = HoodieSparkTable.create(cfg, context, metaClient);
    SparkTaskContextSupplier supplier = new SparkTaskContextSupplier();

    // Should throw an exception when trying to create the writer
    assertThrows(Exception.class, () -> {
      HoodieFileWriterFactory.getFileWriter(
          instantTime, parquetPath, table.getStorage(), cfg.getStorageConfig(),
          schema, supplier, HoodieRecordType.SPARK);
    });
  }

  @Test
  public void testNoInjectorUsesDefaultConfig() throws Exception {
    final String instantTime = "103";
    final StoragePath parquetPath = new StoragePath(
        basePath + "/partition/path/test_no_injector_" + instantTime + ".parquet");

    HoodieSchema schema = getStructTypeCompatibleSchema();

    // Create config WITHOUT injector - should use default settings
    final HoodieWriteConfig cfg = getConfigBuilder()
        .withStorageConfig(HoodieStorageConfig.newBuilder()
            .parquetDictionaryEnabled(true)
            .build())
        .build();

    HoodieTable table = HoodieSparkTable.create(cfg, context, metaClient);
    SparkTaskContextSupplier supplier = new SparkTaskContextSupplier();

    // Create writer and write some data
    HoodieFileWriter writer = HoodieFileWriterFactory.getFileWriter(
        instantTime, parquetPath, table.getStorage(), cfg.getStorageConfig(),
        schema, supplier, HoodieRecordType.SPARK);

    assertTrue(writer instanceof HoodieSparkParquetWriter);

    // Generate test data using SparkDatasetTestUtils and convert to InternalRows
    Dataset<Row> rowDataset = SparkDatasetTestUtils.getRandomRows(sqlContext, 10, "partition/path", false);
    List<InternalRow> internalRows = toMutableRows(SparkDatasetTestUtils.toInternalRows(rowDataset, SparkDatasetTestUtils.ENCODER));
    List<Row> rows = rowDataset.collectAsList();

    // Write some test records
    HoodieSparkParquetWriter sparkWriter = (HoodieSparkParquetWriter) writer;
    for (int i = 0; i < internalRows.size(); i++) {
      InternalRow row = internalRows.get(i);
      String recordKey = rows.get(i).getString(2); // record_key is at index 2 (after commit_time and seq_no)
      String partition = rows.get(i).getString(3);  // partition is at index 3
      sparkWriter.writeRowWithMetadata(new HoodieKey(recordKey, partition), row);
    }
    writer.close();

    // Verify the parquet file was created
    assertTrue(table.getStorage().exists(parquetPath));
  }
}
