/*
 * 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.paimon.operation;

import org.apache.paimon.AppendOnlyFileStore;
import org.apache.paimon.CoreOptions;
import org.apache.paimon.append.AppendOnlyWriter;
import org.apache.paimon.append.cluster.Sorter;
import org.apache.paimon.compact.CompactManager;
import org.apache.paimon.data.BinaryRow;
import org.apache.paimon.data.BlobConsumer;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.deletionvectors.BucketedDvMaintainer;
import org.apache.paimon.deletionvectors.DeletionVector;
import org.apache.paimon.fileindex.FileIndexOptions;
import org.apache.paimon.format.FileFormat;
import org.apache.paimon.fs.FileIO;
import org.apache.paimon.io.BundleRecords;
import org.apache.paimon.io.DataFileMeta;
import org.apache.paimon.io.DataFilePathFactory;
import org.apache.paimon.io.RowDataRollingFileWriter;
import org.apache.paimon.manifest.FileSource;
import org.apache.paimon.metrics.MetricRegistry;
import org.apache.paimon.operation.metrics.BlobFetchMetrics;
import org.apache.paimon.reader.RecordReaderIterator;
import org.apache.paimon.statistics.SimpleColStatsCollector;
import org.apache.paimon.types.RowType;
import org.apache.paimon.utils.CommitIncrement;
import org.apache.paimon.utils.ExceptionUtils;
import org.apache.paimon.utils.FileStorePathFactory;
import org.apache.paimon.utils.IOExceptionSupplier;
import org.apache.paimon.utils.LongCounter;
import org.apache.paimon.utils.MutableObjectIterator;
import org.apache.paimon.utils.RecordWriter;
import org.apache.paimon.utils.SnapshotManager;
import org.apache.paimon.utils.StatsCollectorFactories;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;

import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.function.Function;
import java.util.function.Supplier;

import static org.apache.paimon.format.FileFormat.fileFormat;
import static org.apache.paimon.utils.StatsCollectorFactories.createStatsFactories;

/** {@link FileStoreWrite} for {@link AppendOnlyFileStore}. */
public abstract class BaseAppendFileStoreWrite extends MemoryFileStoreWrite<InternalRow>
        implements BundleFileStoreWriter {

    private static final Logger LOG = LoggerFactory.getLogger(BaseAppendFileStoreWrite.class);

    private final FileIO fileIO;
    private final RawFileSplitRead readForCompact;
    private final long schemaId;
    private final FileFormat fileFormat;
    private final FileStorePathFactory pathFactory;
    private final FileIndexOptions fileIndexOptions;
    private final RowType rowType;

    private @Nullable BlobFileContext blobContext;
    private @Nullable BlobFetchMetrics blobFetchMetrics;
    private RowType writeType;
    private @Nullable List<String> writeCols;
    private boolean forceBufferSpill = false;

    public BaseAppendFileStoreWrite(
            FileIO fileIO,
            RawFileSplitRead readForCompact,
            long schemaId,
            RowType rowType,
            RowType partitionType,
            FileStorePathFactory pathFactory,
            SnapshotManager snapshotManager,
            FileStoreScan scan,
            CoreOptions options,
            @Nullable BucketedDvMaintainer.Factory dvMaintainerFactory,
            String tableName) {
        super(
                snapshotManager,
                scan,
                options,
                partitionType,
                null,
                dvMaintainerFactory,
                null,
                tableName);
        this.fileIO = fileIO;
        this.readForCompact = readForCompact;
        this.schemaId = schemaId;
        this.rowType = rowType;
        this.writeType = rowType;
        this.writeCols = null;
        this.fileFormat = fileFormat(options);
        this.pathFactory = pathFactory;
        this.blobContext = BlobFileContext.create(rowType, options);
        this.fileIndexOptions = options.indexColumnsOptions();
    }

    @Override
    public BaseAppendFileStoreWrite withBlobConsumer(BlobConsumer blobConsumer) {
        if (blobContext != null) {
            blobContext = blobContext.withBlobConsumer(blobConsumer);
        }
        return this;
    }

    @Override
    public BaseAppendFileStoreWrite withMetricRegistry(MetricRegistry metricRegistry) {
        super.withMetricRegistry(metricRegistry);
        if (blobContext != null) {
            blobFetchMetrics = new BlobFetchMetrics(metricRegistry, tableName);
            blobContext = blobContext.withBlobFetchMetricReporter(blobFetchMetrics);
        }
        return this;
    }

    @Override
    protected RecordWriter<InternalRow> createWriter(
            BinaryRow partition,
            int bucket,
            List<DataFileMeta> restoredFiles,
            long restoredMaxSeqNumber,
            @Nullable CommitIncrement restoreIncrement,
            ExecutorService compactExecutor,
            @Nullable BucketedDvMaintainer dvMaintainer,
            boolean ignorePreviousFiles) {
        DataFilePathFactory dataPathFactory =
                pathFactory.createDataFilePathFactory(partition, bucket);
        return new AppendOnlyWriter(
                fileIO,
                ioManager,
                schemaId,
                fileFormat,
                FileFormat.vectorFileFormat(options),
                options.targetFileSize(false),
                options.blobTargetFileSize(),
                options.vectorTargetFileSize(),
                options.targetFileRowNum(),
                writeType,
                writeCols,
                restoredMaxSeqNumber,
                getCompactManager(partition, bucket, restoredFiles, compactExecutor, dvMaintainer),
                // it is only for new files, no dv
                files -> createFilesIterator(partition, bucket, files, null),
                options.commitForceCompact(),
                dataPathFactory,
                restoreIncrement,
                options.useWriteBufferForAppend() || forceBufferSpill,
                options.writeBufferSpillable() || forceBufferSpill,
                options.fileCompression(),
                options.spillCompressOptions(),
                new StatsCollectorFactories(options),
                options.writeBufferSpillDiskSize(),
                fileIndexOptions,
                options.asyncFileWrite(),
                options.statsDenseStore(),
                options.dataEvolutionEnabled(),
                rowSidecarFileFormat(),
                blobContext);
    }

    @Override
    public void withWriteType(RowType writeType) {
        this.writeType = writeType;
        if (blobContext != null) {
            blobContext = blobContext.withWriteType(writeType);
        }
        int fullCount = rowType.getFieldCount();
        List<String> fullNames = rowType.getFieldNames();
        this.writeCols = writeType.getFieldNames();
        // optimize writeCols to null in following cases:
        // writeType contains all columns (without _ROW_ID and _SEQUENCE_NUMBER)
        if (writeCols.equals(fullNames)) {
            writeCols = null;
        }
    }

    private SimpleColStatsCollector.Factory[] statsCollectors() {
        return createStatsFactories(options.statsMode(), options, writeType.getFieldNames());
    }

    @Override
    public void close() throws Exception {
        super.close();
        if (blobFetchMetrics != null) {
            blobFetchMetrics.close();
        }
    }

    protected abstract CompactManager getCompactManager(
            BinaryRow partition,
            int bucket,
            List<DataFileMeta> restoredFiles,
            ExecutorService compactExecutor,
            @Nullable BucketedDvMaintainer dvMaintainer);

    public List<DataFileMeta> compactRewrite(
            BinaryRow partition,
            int bucket,
            @Nullable Function<String, DeletionVector> dvFactory,
            List<DataFileMeta> toCompact)
            throws Exception {
        if (toCompact.isEmpty()) {
            return Collections.emptyList();
        }
        Exception collectedExceptions = null;
        RowDataRollingFileWriter rewriter =
                createRollingFileWriter(
                        partition,
                        bucket,
                        () -> new LongCounter(toCompact.get(0).minSequenceNumber()));
        Map<String, IOExceptionSupplier<DeletionVector>> dvFactories = null;
        if (dvFactory != null) {
            dvFactories = new HashMap<>();
            for (DataFileMeta file : toCompact) {
                dvFactories.put(file.fileName(), () -> dvFactory.apply(file.fileName()));
            }
        }
        try {
            rewriter.write(createFilesIterator(partition, bucket, toCompact, dvFactories));
        } catch (Exception e) {
            collectedExceptions = e;
        } finally {
            try {
                rewriter.close();
            } catch (Exception e) {
                collectedExceptions = ExceptionUtils.firstOrSuppressed(e, collectedExceptions);
            }
        }
        if (collectedExceptions != null) {
            throw collectedExceptions;
        }
        return rewriter.result();
    }

    public List<DataFileMeta> clusterRewrite(
            BinaryRow partition, int bucket, List<DataFileMeta> toCluster) throws Exception {
        RecordReaderIterator<InternalRow> reader =
                createFilesIterator(partition, bucket, toCluster, null);

        // sort and rewrite
        Exception collectedExceptions = null;
        Sorter sorter = Sorter.getSorter(reader, ioManager, rowType, options);
        RowDataRollingFileWriter rewriter =
                createRollingFileWriter(
                        partition,
                        bucket,
                        () -> new LongCounter(toCluster.get(0).minSequenceNumber()));
        try {
            MutableObjectIterator<BinaryRow> sorted = sorter.sort();
            BinaryRow binaryRow = new BinaryRow(sorter.arity());
            while ((binaryRow = sorted.next(binaryRow)) != null) {
                InternalRow rowRemovedKey = sorter.removeSortKey(binaryRow);
                rewriter.write(rowRemovedKey);
            }
        } catch (Exception e) {
            collectedExceptions = e;
        } finally {
            try {
                rewriter.close();
                sorter.close();
            } catch (Exception e) {
                collectedExceptions = ExceptionUtils.firstOrSuppressed(e, collectedExceptions);
            }
        }

        if (collectedExceptions != null) {
            throw collectedExceptions;
        }

        return rewriter.result();
    }

    private RowDataRollingFileWriter createRollingFileWriter(
            BinaryRow partition, int bucket, Supplier<LongCounter> seqNumCounterSupplier) {
        return new RowDataRollingFileWriter(
                fileIO,
                schemaId,
                fileFormat,
                options.targetFileSize(false),
                writeType,
                pathFactory.createDataFilePathFactory(partition, bucket),
                seqNumCounterSupplier,
                options.fileCompression(),
                statsCollectors(),
                fileIndexOptions,
                FileSource.COMPACT,
                options.asyncFileWrite(),
                options.statsDenseStore(),
                rowType.equals(writeType) ? null : writeType.getFieldNames(),
                rowSidecarFileFormat(),
                Long.MAX_VALUE);
    }

    @Nullable
    private FileFormat rowSidecarFileFormat() {
        return options.dataEvolutionEnabled() && options.dataEvolutionRowSidecarEnabled()
                ? FileFormat.fromIdentifier("row", options.toConfiguration())
                : null;
    }

    private RecordReaderIterator<InternalRow> createFilesIterator(
            BinaryRow partition,
            int bucket,
            List<DataFileMeta> files,
            @Nullable Map<String, IOExceptionSupplier<DeletionVector>> dvFactories)
            throws IOException {
        return new RecordReaderIterator<>(
                readForCompact.createReader(partition, bucket, files, dvFactories));
    }

    @Override
    protected void forceBufferSpill() throws Exception {
        if (ioManager == null) {
            return;
        }
        if (blobContext != null) {
            return;
        }
        if (forceBufferSpill) {
            return;
        }
        forceBufferSpill = true;
        LOG.info(
                "Force buffer spill for append-only file store write, writer number is: {}",
                writers.size());
        for (Map<Integer, WriterContainer<InternalRow>> bucketWriters : writers.values()) {
            for (WriterContainer<InternalRow> writerContainer : bucketWriters.values()) {
                ((AppendOnlyWriter) writerContainer.writer).toBufferedWriter();
            }
        }
    }

    @Override
    public void writeBundle(BinaryRow partition, int bucket, BundleRecords bundle)
            throws Exception {
        WriterContainer<InternalRow> container = getWriterWrapper(partition, bucket);
        ((AppendOnlyWriter) container.writer).writeBundle(bundle);
    }
}
