/*
 * 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.flink.sink;

import java.util.Collection;
import java.util.Map;
import java.util.Set;
import java.util.function.Supplier;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.types.logical.RowType;
import org.apache.iceberg.FileFormat;
import org.apache.iceberg.PartitionKey;
import org.apache.iceberg.PartitionSpec;
import org.apache.iceberg.Schema;
import org.apache.iceberg.Table;
import org.apache.iceberg.TableUtil;
import org.apache.iceberg.flink.RowDataWrapper;
import org.apache.iceberg.io.FileIO;
import org.apache.iceberg.io.FileWriterFactory;
import org.apache.iceberg.io.OutputFileFactory;
import org.apache.iceberg.io.PartitionedFanoutWriter;
import org.apache.iceberg.io.TaskWriter;
import org.apache.iceberg.io.UnpartitionedWriter;
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
import org.apache.iceberg.relocated.com.google.common.collect.Sets;
import org.apache.iceberg.types.TypeUtil;
import org.apache.iceberg.util.ArrayUtil;
import org.apache.iceberg.util.SerializableSupplier;

public class RowDataTaskWriterFactory implements TaskWriterFactory<RowData> {
  private final Supplier<Table> tableSupplier;
  private final Schema schema;
  private final RowType flinkSchema;
  private final PartitionSpec spec;
  private final long targetFileSizeBytes;
  private final FileFormat format;
  private final Set<Integer> equalityFieldIds;
  private final boolean upsert;
  private final FileWriterFactory<RowData> fileWriterFactory;
  private boolean useDv;

  private transient OutputFileFactory outputFileFactory;

  public RowDataTaskWriterFactory(
      Table table,
      RowType flinkSchema,
      long targetFileSizeBytes,
      FileFormat format,
      Map<String, String> writeProperties,
      Collection<Integer> equalityFieldIds,
      boolean upsert) {
    this(
        () -> table,
        flinkSchema,
        targetFileSizeBytes,
        format,
        writeProperties,
        equalityFieldIds,
        upsert);
  }

  public RowDataTaskWriterFactory(
      SerializableSupplier<Table> tableSupplier,
      RowType flinkSchema,
      long targetFileSizeBytes,
      FileFormat format,
      Map<String, String> writeProperties,
      Collection<Integer> equalityFieldIds,
      boolean upsert) {
    this(
        tableSupplier,
        flinkSchema,
        targetFileSizeBytes,
        format,
        writeProperties,
        equalityFieldIds,
        upsert,
        tableSupplier.get().schema(),
        tableSupplier.get().spec());
  }

  public RowDataTaskWriterFactory(
      SerializableSupplier<Table> tableSupplier,
      RowType flinkSchema,
      long targetFileSizeBytes,
      FileFormat format,
      Map<String, String> writeProperties,
      Collection<Integer> equalityFieldIds,
      boolean upsert,
      Schema schema,
      PartitionSpec spec) {
    this.tableSupplier = tableSupplier;

    Table table;
    if (tableSupplier instanceof CachingTableSupplier) {
      // rely on the initial table metadata for schema, etc., until schema evolution is supported
      table = ((CachingTableSupplier) tableSupplier).initialTable();
    } else {
      table = tableSupplier.get();
    }

    this.schema = schema;
    this.flinkSchema = flinkSchema;
    this.spec = spec;
    this.targetFileSizeBytes = targetFileSizeBytes;
    this.format = format;
    this.equalityFieldIds = equalityFieldIds != null ? Sets.newHashSet(equalityFieldIds) : null;
    this.upsert = upsert;

    if (equalityFieldIds == null || equalityFieldIds.isEmpty()) {
      this.fileWriterFactory =
          new FlinkFileWriterFactory.Builder(table)
              .dataFileFormat(format)
              .dataSchema(schema)
              .dataFlinkType(flinkSchema)
              .writerProperties(writeProperties)
              .build();
    } else if (upsert) {
      // In upsert mode, only the new row is emitted using INSERT row kind. Therefore, any column of
      // the inserted row
      // may differ from the deleted row other than the primary key fields, and the delete file must
      // contain values
      // that are correct for the deleted row. Therefore, only write the equality delete fields.
      this.fileWriterFactory =
          new FlinkFileWriterFactory.Builder(table)
              .dataFileFormat(format)
              .dataSchema(schema)
              .dataFlinkType(flinkSchema)
              .deleteFileFormat(format)
              .equalityFieldIds(ArrayUtil.toPrimitive(equalityFieldIds.toArray(new Integer[0])))
              .equalityDeleteRowSchema(TypeUtil.select(schema, Sets.newHashSet(equalityFieldIds)))
              .writerProperties(writeProperties)
              .build();
    } else {
      this.fileWriterFactory =
          new FlinkFileWriterFactory.Builder(table)
              .dataFileFormat(format)
              .dataSchema(schema)
              .dataFlinkType(flinkSchema)
              .deleteFileFormat(format)
              .equalityFieldIds(ArrayUtil.toPrimitive(equalityFieldIds.toArray(new Integer[0])))
              .equalityDeleteRowSchema(schema)
              .writerProperties(writeProperties)
              .build();
    }
  }

  @Override
  public void initialize(int taskId, int attemptId) {
    Table table;
    if (tableSupplier instanceof CachingTableSupplier) {
      // rely on the initial table metadata for schema, etc., until schema evolution is supported
      table = ((CachingTableSupplier) tableSupplier).initialTable();
    } else {
      table = tableSupplier.get();
    }

    refreshTable();
    this.useDv = TableUtil.formatVersion(table) > 2;

    this.outputFileFactory =
        OutputFileFactory.builderFor(table, taskId, attemptId)
            .format(format)
            .ioSupplier(() -> tableSupplier.get().io())
            .defaultSpec(spec)
            .build();
  }

  @Override
  public TaskWriter<RowData> create() {
    Preconditions.checkNotNull(
        outputFileFactory,
        "The outputFileFactory shouldn't be null if we have invoked the initialize().");

    refreshTable();

    if (equalityFieldIds == null || equalityFieldIds.isEmpty()) {
      // Initialize a task writer to write INSERT only.
      if (spec.isUnpartitioned()) {
        return new UnpartitionedWriter<>(
            spec,
            format,
            fileWriterFactory,
            outputFileFactory,
            tableSupplier.get().io(),
            targetFileSizeBytes);
      } else {
        return new RowDataPartitionedFanoutWriter(
            spec,
            format,
            fileWriterFactory,
            outputFileFactory,
            tableSupplier.get().io(),
            targetFileSizeBytes,
            schema,
            flinkSchema);
      }
    } else {
      // Initialize a task writer to write both INSERT and equality DELETE.
      if (spec.isUnpartitioned()) {
        return new UnpartitionedDeltaWriter(
            spec,
            format,
            fileWriterFactory,
            outputFileFactory,
            tableSupplier.get().io(),
            targetFileSizeBytes,
            schema,
            flinkSchema,
            equalityFieldIds,
            upsert,
            useDv);
      } else {
        return new PartitionedDeltaWriter(
            spec,
            format,
            fileWriterFactory,
            outputFileFactory,
            tableSupplier.get().io(),
            targetFileSizeBytes,
            schema,
            flinkSchema,
            equalityFieldIds,
            upsert,
            useDv);
      }
    }
  }

  void refreshTable() {
    if (tableSupplier instanceof CachingTableSupplier) {
      ((CachingTableSupplier) tableSupplier).refreshTable();
    }
  }

  private static class RowDataPartitionedFanoutWriter extends PartitionedFanoutWriter<RowData> {

    private final PartitionKey partitionKey;
    private final RowDataWrapper rowDataWrapper;

    RowDataPartitionedFanoutWriter(
        PartitionSpec spec,
        FileFormat format,
        FileWriterFactory<RowData> fileWriterFactory,
        OutputFileFactory fileFactory,
        FileIO io,
        long targetFileSize,
        Schema schema,
        RowType flinkSchema) {
      super(spec, format, fileWriterFactory, fileFactory, io, targetFileSize);
      this.partitionKey = new PartitionKey(spec, schema);
      this.rowDataWrapper = new RowDataWrapper(flinkSchema, schema.asStruct());
    }

    @Override
    protected PartitionKey partition(RowData row) {
      partitionKey.partition(rowDataWrapper.wrap(row));
      return partitionKey;
    }
  }
}
