/*
 * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
 */

package io.airbyte.integrations.destination.record_buffer;

import com.google.common.io.CountingOutputStream;
import io.airbyte.protocol.models.AirbyteRecordMessage;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import org.apache.commons.compress.compressors.gzip.GzipCompressorOutputStream;
import org.apache.commons.io.FileUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Base implementation of a {@link SerializableBuffer}. It is composed of a {@link BufferStorage}
 * where the actual data is being stored in a serialized format.
 *
 * Such data format is defined by concrete implementation inheriting from this base abstract class.
 * To do so, necessary methods on handling "writer" methods should be defined. This writer would
 * take care of converting {@link AirbyteRecordMessage} into the serialized form of the data such as
 * it can be stored in the outputStream of the {@link BufferStorage}.
 */
public abstract class BaseSerializedBuffer implements SerializableBuffer {

  private static final Logger LOGGER = LoggerFactory.getLogger(BaseSerializedBuffer.class);
  private static final String GZ_SUFFIX = ".gz";

  private final BufferStorage bufferStorage;
  private final CountingOutputStream byteCounter;

  private boolean useCompression;
  private GzipCompressorOutputStream compressedBuffer;
  private InputStream inputStream;
  private boolean isStarted;
  private boolean isClosed;

  protected BaseSerializedBuffer(final BufferStorage bufferStorage) throws Exception {
    this.bufferStorage = bufferStorage;
    byteCounter = new CountingOutputStream(bufferStorage.getOutputStream());
    useCompression = true;
    compressedBuffer = null;
    inputStream = null;
    isStarted = false;
    isClosed = false;
  }

  /**
   * Initializes the writer objects such that it can now write to the downstream @param outputStream
   */
  protected abstract void createWriter(OutputStream outputStream) throws Exception;

  /**
   * Transform the @param recordMessage into a serialized form of the data and writes it to the
   * registered OutputStream provided when {@link BaseSerializedBuffer#createWriter} was called.
   */
  protected abstract void writeRecord(AirbyteRecordMessage recordMessage) throws IOException;

  /**
   * Stops the writer from receiving new data and prepares it for being finalized and converted into
   * an InputStream to read from instead. This is used when flushing the buffer into some other
   * destination.
   */
  protected abstract void flushWriter() throws IOException;

  protected abstract void closeWriter() throws IOException;

  public SerializableBuffer withCompression(final boolean useCompression) {
    if (!isStarted) {
      this.useCompression = useCompression;
      return this;
    }
    throw new RuntimeException("Options should be configured before starting to write");
  }

  @Override
  public long accept(final AirbyteRecordMessage recordMessage) throws Exception {
    if (!isStarted) {
      if (useCompression) {
        compressedBuffer = new GzipCompressorOutputStream(byteCounter);
        createWriter(compressedBuffer);
      } else {
        createWriter(byteCounter);
      }
      isStarted = true;
    }
    if (inputStream == null && !isClosed) {
      final long startCount = byteCounter.getCount();
      writeRecord(recordMessage);
      return byteCounter.getCount() - startCount;
    } else {
      throw new IllegalCallerException("Buffer is already closed, it cannot accept more messages");
    }
  }

  @Override
  public String getFilename() throws IOException {
    if (useCompression && !bufferStorage.getFilename().endsWith(GZ_SUFFIX)) {
      return bufferStorage.getFilename() + GZ_SUFFIX;
    }
    return bufferStorage.getFilename();
  }

  @Override
  public File getFile() throws IOException {
    if (useCompression && !bufferStorage.getFilename().endsWith(GZ_SUFFIX)) {
      if (bufferStorage.getFile().renameTo(new File(bufferStorage.getFilename() + GZ_SUFFIX))) {
        LOGGER.info("Renaming compressed file to include .gz file extension");
      }
    }
    return bufferStorage.getFile();
  }

  protected InputStream convertToInputStream() throws IOException {
    return bufferStorage.convertToInputStream();
  }

  @Override
  public InputStream getInputStream() {
    return inputStream;
  }

  @Override
  public void flush() throws IOException {
    if (inputStream == null && !isClosed) {
      flushWriter();
      if (compressedBuffer != null) {
        LOGGER.info("Wrapping up compression and write GZIP trailer data.");
        compressedBuffer.flush();
        compressedBuffer.close();
      }
      closeWriter();
      bufferStorage.close();
      inputStream = convertToInputStream();
      LOGGER.info("Finished writing data to {} ({})", getFilename(), FileUtils.byteCountToDisplaySize(byteCounter.getCount()));
    }
  }

  @Override
  public long getByteCount() {
    return byteCounter.getCount();
  }

  @Override
  public void close() throws Exception {
    if (!isClosed) {
      // inputStream can be null if the accept method encounters
      // an error before inputStream is initialized
      if (inputStream != null) {
        inputStream.close();
      }
      bufferStorage.deleteFile();
      isClosed = true;
    }
  }

  @Override
  public long getMaxTotalBufferSizeInBytes() {
    return bufferStorage.getMaxTotalBufferSizeInBytes();
  }

  @Override
  public long getMaxPerStreamBufferSizeInBytes() {
    return bufferStorage.getMaxPerStreamBufferSizeInBytes();
  }

  @Override
  public int getMaxConcurrentStreamsInBuffer() {
    return bufferStorage.getMaxConcurrentStreamsInBuffer();
  }

}
