package com.atlassian.bitbucket.mesh.git.hook;

import com.google.common.io.ByteStreams;
import com.google.protobuf.ByteString;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.BufferOverflowException;
import java.nio.ByteBuffer;
import java.nio.channels.WritableByteChannel;
import java.nio.file.Files;
import java.nio.file.NoSuchFileException;
import java.nio.file.Path;

import static java.nio.file.StandardOpenOption.APPEND;
import static java.nio.file.StandardOpenOption.WRITE;
import static java.util.Objects.requireNonNull;

class HookStdinBuilder {

    private final long bufferCapacity;
    private final Path tempDir;

    private ByteString buffered;
    private long size;
    private Path written;

    HookStdinBuilder(@Nullable Path tempDir, long bufferCapacity) {
        this.bufferCapacity = bufferCapacity;
        this.tempDir = tempDir;

        buffered = ByteString.EMPTY;
    }

    void append(@Nonnull ByteBuffer buffer) throws IOException {
        if (!requireNonNull(buffer, "buffer").hasRemaining()) {
            return;
        }

        try {
            size += doAppend(buffer);
        } catch (IOException | RuntimeException e) {
            //If any exception is thrown while we're still building we need to ensure we clean up
            //the temporary file if we've created one
            if (written != null) {
                try {
                    Files.delete(written);
                    written = null;
                } catch (FileNotFoundException | NoSuchFileException ignored) {
                    //If the file doesn't exist, don't suppress the exception
                } catch (IOException suppressed) {
                    e.addSuppressed(suppressed);
                }
            }

            throw e;
        }
    }

    @Nonnull
    HookStdin build() {
        if (written == null) {
            return new ByteStringHookStdin(buffered, size);
        }

        return new PathHookStdin(written, size);
    }

    private int doAppend(ByteBuffer buffer) throws IOException {
        int remaining = buffer.remaining();
        if (written == null) {
            if (size + remaining <= bufferCapacity) {
                //The input will still fit in memory, so append it to our buffer
                ByteString chunk = ByteString.copyFrom(buffer);
                if (size == 0L) {
                    buffered = chunk;
                } else {
                    buffered = buffered.concat(chunk);
                }

                return chunk.size();
            }
            if (tempDir == null) {
                //We're not allowed to buffer to disk, so throw instead
                throw new BufferOverflowException();
            }

            //We've received more stdin than we can buffer in memory, so we need to switch to
            //buffering to disk
            written = Files.createTempFile(tempDir, "hook-stdin", ".dat");
            if (size > 0L) {
                //Don't use ByteString.writeTo; it copies all byte[]s before providing them to
                //OutputStream.write. newInput + copy avoids all those copies and only creates
                //a single (reused) byte[] to copy through
                try (InputStream inputStream = buffered.newInput();
                     OutputStream outputStream = Files.newOutputStream(written)) {
                    ByteStreams.copy(inputStream, outputStream);
                }
            }
            buffered = null; //Release any buffered memory (this may be ByteString.EMPTY)
        }

        try (WritableByteChannel channel = Files.newByteChannel(written, APPEND, WRITE)) {
            //The size returned here intentionally does not include any buffered bytes written;
            //they've already been counted on previous append calls
            return channel.write(buffer);
        }
    }
}
