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

import org.apache.paimon.data.InternalRow;
import org.apache.paimon.format.BundleFormatWriter;
import org.apache.paimon.format.FileAwareFormatWriter;
import org.apache.paimon.format.FormatWriter;
import org.apache.paimon.format.FormatWriterFactory;
import org.apache.paimon.format.SupportsDirectWrite;
import org.apache.paimon.fs.AsyncPositionOutputStream;
import org.apache.paimon.fs.FileIO;
import org.apache.paimon.fs.Path;
import org.apache.paimon.fs.PositionOutputStream;
import org.apache.paimon.utils.IOUtils;

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

import javax.annotation.Nullable;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.util.Optional;
import java.util.function.Function;

/**
 * A {@link FileWriter} to produce a single file.
 *
 * @param <T> type of records to write.
 * @param <R> type of result to produce after writing a file.
 */
public abstract class SingleFileWriter<T, R> implements FileWriter<T, R> {

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

    protected final FileIO fileIO;
    protected final Path path;
    private final Function<T, InternalRow> converter;

    private boolean deleteFileUponAbort;
    private FormatWriter writer;
    @Nullable private PositionOutputStream out;

    @Nullable private Long outputBytes;
    @Nullable private Object writerMetadata;
    private long recordCount;
    protected boolean closed;

    public SingleFileWriter(
            FileIO fileIO,
            FormatWriterFactory factory,
            Path path,
            Function<T, InternalRow> converter,
            String compression,
            boolean asyncWrite) {
        this.fileIO = fileIO;
        this.path = path;
        this.converter = converter;
        // true first to clean file in exception
        this.deleteFileUponAbort = true;

        try {
            if (factory instanceof SupportsDirectWrite) {
                writer = ((SupportsDirectWrite) factory).create(fileIO, path, compression);
            } else {
                out = fileIO.newOutputStream(path, false);
                if (asyncWrite) {
                    out = new AsyncPositionOutputStream(out);
                }
                writer = factory.create(out, compression);
            }

            if (writer instanceof FileAwareFormatWriter) {
                FileAwareFormatWriter fileAwareFormatWriter = (FileAwareFormatWriter) writer;
                fileAwareFormatWriter.setFile(path);
                deleteFileUponAbort = fileAwareFormatWriter.deleteFileUponAbort();
            }
        } catch (IOException e) {
            LOG.warn(
                    "Failed to open the bulk writer, closing the output stream and throw the error.",
                    e);
            if (out != null) {
                abort();
            }
            throw new UncheckedIOException(e);
        }

        this.recordCount = 0;
        this.closed = false;
    }

    public Path path() {
        return path;
    }

    @Override
    public void write(T record) throws IOException {
        writeImpl(record);
    }

    public void writeBundle(BundleRecords bundle) throws IOException {
        if (closed) {
            throw new RuntimeException("Writer has already closed!");
        }

        try {
            if (writer instanceof BundleFormatWriter) {
                ((BundleFormatWriter) writer).writeBundle(bundle);
            } else {
                for (InternalRow row : bundle) {
                    writer.addElement(row);
                }
            }
            recordCount += bundle.rowCount();
        } catch (Throwable e) {
            LOG.warn("Exception occurs when writing file {}. Cleaning up.", path, e);
            abort();
            throw e;
        }
    }

    protected InternalRow writeImpl(T record) throws IOException {
        if (closed) {
            throw new RuntimeException("Writer has already closed!");
        }

        try {
            InternalRow rowData = converter.apply(record);
            writer.addElement(rowData);
            recordCount++;
            return rowData;
        } catch (Throwable e) {
            LOG.warn("Exception occurs when writing file {}. Cleaning up.", path, e);
            abort();
            throw e;
        }
    }

    @Override
    public long recordCount() {
        return recordCount;
    }

    public boolean reachTargetSize(boolean suggestedCheck, long targetSize) throws IOException {
        return writer.reachTargetSize(suggestedCheck, targetSize);
    }

    @Override
    public void abort() {
        if (writer != null) {
            IOUtils.closeQuietly(writer);
            writer = null;
        }
        if (out != null) {
            IOUtils.closeQuietly(out);
            out = null;
        }
        abortExecutor().ifPresent(FileWriterAbortExecutor::abort);
    }

    public Optional<FileWriterAbortExecutor> abortExecutor() {
        return deleteFileUponAbort
                ? Optional.of(new FileWriterAbortExecutor(fileIO, path))
                : Optional.empty();
    }

    @Override
    public void close() throws IOException {
        if (closed) {
            return;
        }

        if (LOG.isDebugEnabled()) {
            LOG.debug("Closing file {}", path);
        }

        try {
            if (writer != null) {
                writer.close();
                writerMetadata = writer.writerMetadata();
                writer = null;
            }
            if (out != null) {
                out.flush();
                outputBytes = out.getPos();
                out.close();
                out = null;
            }
        } catch (IOException e) {
            LOG.warn("Exception occurs when closing file {}. Cleaning up.", path, e);
            abort();
            throw e;
        } finally {
            closed = true;
        }
    }

    protected long outputBytes() throws IOException {
        if (outputBytes == null) {
            outputBytes = fileIO.getFileSize(path);
        }
        return outputBytes;
    }

    /**
     * Returns cached writer metadata from the format writer. Available after {@link #close()} is
     * called. Can be used by stats extractors to avoid re-reading the file from object storage.
     */
    @Nullable
    protected Object writerMetadata() {
        return writerMetadata;
    }
}
