coding-agent-benchmark/src/main/java/illgner/ch/createzip/SplitZipWriter.java

package illgner.ch.createzip;

import illgner.ch.stream.OutputStreamEnhancer;
import org.apache.commons.io.IOUtils;
import org.apache.commons.io.output.CountingOutputStream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.crypto.NoSuchPaddingException;
import java.io.BufferedOutputStream;
import java.io.Closeable;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.attribute.PosixFilePermission;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.util.EnumSet;
import java.util.Set;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

import static java.nio.file.attribute.PosixFilePermission.GROUP_READ;
import static java.nio.file.attribute.PosixFilePermission.OTHERS_READ;
import static java.nio.file.attribute.PosixFilePermission.OWNER_READ;
import static java.nio.file.attribute.PosixFilePermission.OWNER_WRITE;

public class SplitZipWriter implements Closeable {

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

    private static final Set<PosixFilePermission> ZIP_FILE_PERMISSIONS = EnumSet.of(OWNER_READ, OWNER_WRITE, GROUP_READ, OTHERS_READ);

    private final Path destinationDir;
    private final String baseName;
    private final int chunkSize;
    private final int bufferSize;
    private final ZipResult zipResult;
    private OutputStreamEnhancer outputStreamEnhancer;

    private ZipOutputStream zipOutputStream;

    private int fileCount = 1;
    private long bytesWritten = 0;
    private File actualZipFile;
    private CountingOutputStream countingOutputStream;

    public SplitZipWriter(Path destinationDir, String baseName, int chunkSize, int bufferSize, OutputStreamEnhancer outputStreamEnhancer) {
        this.destinationDir = destinationDir;
        this.baseName = baseName;
        this.chunkSize = chunkSize;
        this.bufferSize = bufferSize;
        this.zipResult = new ZipResult();
        this.outputStreamEnhancer = outputStreamEnhancer;
    }

    private ZipOutputStream createNewZipOutputStream() throws IOException, InvalidAlgorithmParameterException, NoSuchPaddingException,
            NoSuchAlgorithmException, InvalidKeyException, NoSuchProviderException {

        String actualZipFilename = createOutputFilename();
        actualZipFile = destinationDir.resolve(actualZipFilename).toFile();

        LOG.info("Creating zip part {}: {}", fileCount, actualZipFilename);

        countingOutputStream = new CountingOutputStream(outputStreamEnhancer.enhance(new BufferedOutputStream(new FileOutputStream(actualZipFile), bufferSize)));
        return new ZipOutputStream(countingOutputStream);
    }

    private String createOutputFilename() {
        return String.format("%s_%02d.zip", baseName, fileCount);
    }

    public void write(InputStream inputStream, String path) throws IOException, InvalidAlgorithmParameterException, NoSuchPaddingException,
            NoSuchAlgorithmException, InvalidKeyException, NoSuchProviderException {

        if (countingOutputStream == null) {
            zipOutputStream = createNewZipOutputStream();
        }
        if (bytesWritten >= chunkSize) {
            switchOutputFile();
        }
        LOG.info("Writing file {} to zip", path);
        zipOutputStream.putNextEntry(new ZipEntry(path));
        bytesWritten += IOUtils.copyLarge(inputStream, zipOutputStream);
        zipOutputStream.closeEntry();
    }

    private void switchOutputFile() throws IOException, InvalidAlgorithmParameterException, NoSuchPaddingException, NoSuchAlgorithmException, InvalidKeyException, NoSuchProviderException {
        LOG.info("Switching zip");
        zipOutputStream.flush();
        zipOutputStream.close();
        addZipToResult();
        fileCount++;
        zipOutputStream = createNewZipOutputStream();
        bytesWritten = 0;
    }

    public void close() throws IOException {
        LOG.info("Closing zip {}", actualZipFile.getName());
        zipOutputStream.flush();
        zipOutputStream.close();
        addZipToResult();
    }

    public ZipResult getZipResult() {
        return this.zipResult;
    }

    private void addZipToResult() throws IOException {
        zipResult.add(actualZipFile.getName(), countingOutputStream.getCount());
        Files.setPosixFilePermissions(actualZipFile.toPath(), ZIP_FILE_PERMISSIONS);
    }
}