/*
 * 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.iceberg.parquet;

import org.apache.amoro.shade.guava32.com.google.common.collect.ImmutableList;
import org.apache.amoro.shade.guava32.com.google.common.collect.ImmutableMap;
import org.apache.iceberg.MetadataColumns;
import org.apache.iceberg.Schema;
import org.apache.iceberg.exceptions.RuntimeIOException;
import org.apache.iceberg.expressions.Expression;
import org.apache.iceberg.io.InputFile;
import org.apache.iceberg.mapping.NameMapping;
import org.apache.parquet.ParquetReadOptions;
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.ColumnPath;
import org.apache.parquet.schema.MessageType;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * Copy from iceberg {@ling org.apache.iceberg.parquet.ReadConf} to use {@link
 * AdaptHiveParquetMetricsRowGroupFilter} and {@link AdaptHiveParquetDictionaryRowGroupFilter}
 */
class AdaptHiveReadConf<T> {
  private final ParquetFileReader reader;
  private final InputFile file;
  private final ParquetReadOptions options;
  private final MessageType projection;
  private final ParquetValueReader<T> model;
  private final VectorizedReader<T> vectorizedModel;
  private final List<BlockMetaData> rowGroups;
  private final boolean[] shouldSkip;
  private final long totalValues;
  private final boolean reuseContainers;
  private final Integer batchSize;
  private final long[] startRowPositions;

  // List of column chunk metadata for each row group
  private final List<Map<ColumnPath, ColumnChunkMetaData>> columnChunkMetaDataForRowGroups;

  @SuppressWarnings("unchecked")
  AdaptHiveReadConf(
      InputFile file,
      ParquetReadOptions options,
      Schema expectedSchema,
      Expression filter,
      Function<MessageType, ParquetValueReader<?>> readerFunc,
      Function<MessageType, VectorizedReader<?>> batchedReaderFunc,
      NameMapping nameMapping,
      boolean reuseContainers,
      boolean caseSensitive,
      Integer batchSize) {
    this.file = file;
    this.options = options;
    this.reader = newReader(file, options);
    MessageType fileSchema = reader.getFileMetaData().getSchema();

    MessageType typeWithIds;
    if (ParquetSchemaUtil.hasIds(fileSchema)) {
      typeWithIds = fileSchema;
      this.projection = ParquetSchemaUtil.pruneColumns(fileSchema, expectedSchema);
    } else if (nameMapping != null) {
      typeWithIds =
          (MessageType)
              ParquetTypeVisitor.visit(fileSchema, new AdaptHiveApplyNameMapping(nameMapping));
      this.projection = ParquetSchemaUtil.pruneColumns(typeWithIds, expectedSchema);
    } else {
      typeWithIds = ParquetSchemaUtil.addFallbackIds(fileSchema);
      this.projection = ParquetSchemaUtil.pruneColumnsFallback(fileSchema, expectedSchema);
    }

    this.rowGroups = reader.getRowGroups();
    this.shouldSkip = new boolean[rowGroups.size()];
    this.startRowPositions = new long[rowGroups.size()];

    // Fetch all row groups starting positions to compute the row offsets of the filtered row groups
    Map<Long, Long> offsetToStartPos = generateOffsetToStartPos(expectedSchema);

    // Change for mixed-hive table ⬇
    // Use mixed-hive format filter
    AdaptHiveParquetMetricsRowGroupFilter statsFilter = null;
    AdaptHiveParquetDictionaryRowGroupFilter dictFilter = null;
    if (filter != null) {
      statsFilter =
          new AdaptHiveParquetMetricsRowGroupFilter(expectedSchema, filter, caseSensitive);
      dictFilter =
          new AdaptHiveParquetDictionaryRowGroupFilter(expectedSchema, filter, caseSensitive);
    }
    // Change for mixed-hive table ⬆

    long computedTotalValues = 0L;
    for (int i = 0; i < shouldSkip.length; i += 1) {
      BlockMetaData rowGroup = rowGroups.get(i);
      startRowPositions[i] =
          offsetToStartPos == null ? 0 : offsetToStartPos.get(rowGroup.getStartingPos());
      boolean shouldRead =
          filter == null
              || (statsFilter.shouldRead(typeWithIds, rowGroup)
                  && dictFilter.shouldRead(
                      typeWithIds, rowGroup, reader.getDictionaryReader(rowGroup)));
      this.shouldSkip[i] = !shouldRead;
      if (shouldRead) {
        computedTotalValues += rowGroup.getRowCount();
      }
    }

    this.totalValues = computedTotalValues;
    if (readerFunc != null) {
      this.model = (ParquetValueReader<T>) readerFunc.apply(typeWithIds);
      this.vectorizedModel = null;
      this.columnChunkMetaDataForRowGroups = null;
    } else {
      this.model = null;
      this.vectorizedModel = (VectorizedReader<T>) batchedReaderFunc.apply(typeWithIds);
      this.columnChunkMetaDataForRowGroups = getColumnChunkMetadataForRowGroups();
    }

    this.reuseContainers = reuseContainers;
    this.batchSize = batchSize;
  }

  private AdaptHiveReadConf(AdaptHiveReadConf<T> toCopy) {
    this.reader = null;
    this.file = toCopy.file;
    this.options = toCopy.options;
    this.projection = toCopy.projection;
    this.model = toCopy.model;
    this.rowGroups = toCopy.rowGroups;
    this.shouldSkip = toCopy.shouldSkip;
    this.totalValues = toCopy.totalValues;
    this.reuseContainers = toCopy.reuseContainers;
    this.batchSize = toCopy.batchSize;
    this.vectorizedModel = toCopy.vectorizedModel;
    this.columnChunkMetaDataForRowGroups = toCopy.columnChunkMetaDataForRowGroups;
    this.startRowPositions = toCopy.startRowPositions;
  }

  ParquetFileReader reader() {
    if (reader != null) {
      reader.setRequestedSchema(projection);
      return reader;
    }

    ParquetFileReader newReader = newReader(file, options);
    newReader.setRequestedSchema(projection);
    return newReader;
  }

  ParquetValueReader<T> model() {
    return model;
  }

  VectorizedReader<T> vectorizedModel() {
    return vectorizedModel;
  }

  boolean[] shouldSkip() {
    return shouldSkip;
  }

  private Map<Long, Long> generateOffsetToStartPos(Schema schema) {
    if (schema.findField(MetadataColumns.ROW_POSITION.fieldId()) == null) {
      return null;
    }

    try (ParquetFileReader fileReader = newReader(file, ParquetReadOptions.builder().build())) {
      Map<Long, Long> offsetToStartPos = new HashMap<>();

      long curRowCount = 0;
      for (int i = 0; i < fileReader.getRowGroups().size(); i += 1) {
        BlockMetaData meta = fileReader.getRowGroups().get(i);
        offsetToStartPos.put(meta.getStartingPos(), curRowCount);
        curRowCount += meta.getRowCount();
      }

      return offsetToStartPos;

    } catch (IOException e) {
      throw new UncheckedIOException("Failed to create/close reader for file: " + file, e);
    }
  }

  long[] startRowPositions() {
    return startRowPositions;
  }

  long totalValues() {
    return totalValues;
  }

  boolean reuseContainers() {
    return reuseContainers;
  }

  Integer batchSize() {
    return batchSize;
  }

  List<Map<ColumnPath, ColumnChunkMetaData>> columnChunkMetadataForRowGroups() {
    return columnChunkMetaDataForRowGroups;
  }

  AdaptHiveReadConf<T> copy() {
    return new AdaptHiveReadConf<>(this);
  }

  private static ParquetFileReader newReader(InputFile file, ParquetReadOptions options) {
    try {
      return ParquetFileReader.open(ParquetIO.file(file), options);
    } catch (IOException e) {
      throw new RuntimeIOException(e, "Failed to open Parquet file: %s", file.location());
    }
  }

  private List<Map<ColumnPath, ColumnChunkMetaData>> getColumnChunkMetadataForRowGroups() {
    Set<ColumnPath> projectedColumns =
        projection.getColumns().stream()
            .map(columnDescriptor -> ColumnPath.get(columnDescriptor.getPath()))
            .collect(Collectors.toSet());
    ImmutableList.Builder<Map<ColumnPath, ColumnChunkMetaData>> listBuilder =
        ImmutableList.builder();
    for (int i = 0; i < rowGroups.size(); i++) {
      if (!shouldSkip[i]) {
        BlockMetaData blockMetaData = rowGroups.get(i);
        ImmutableMap.Builder<ColumnPath, ColumnChunkMetaData> mapBuilder = ImmutableMap.builder();
        blockMetaData.getColumns().stream()
            .filter(columnChunkMetaData -> projectedColumns.contains(columnChunkMetaData.getPath()))
            .forEach(
                columnChunkMetaData ->
                    mapBuilder.put(columnChunkMetaData.getPath(), columnChunkMetaData));
        listBuilder.add(mapBuilder.build());
      } else {
        listBuilder.add(ImmutableMap.of());
      }
    }
    return listBuilder.build();
  }
}
