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

import com.atlassian.bitbucket.mesh.git.hosting.HostingOperationType;
import com.atlassian.util.contentcache.CacheResult;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.channels.Channels;
import java.nio.channels.WritableByteChannel;
import java.nio.charset.CharsetEncoder;
import java.nio.charset.CoderResult;
import java.nio.charset.CodingErrorAction;
import java.nio.charset.StandardCharsets;

import static com.google.common.base.Preconditions.checkArgument;
import static java.util.Objects.requireNonNull;

/**
 * Writes Git hook output to a channel.
 * <p>
 * Messages written to {@code stderr} or {@code stdout} are constructed as a series of chunks in the following format:
 * <code>[ASCII char][unsigned short][bytes]</code>
 * The meaning of each part of the chunk is:
 * <dl>
 *     <dt>ASCII char</dt>
 *     <dd>The <em>channel</em> of the data. Can be standard output (O) or standard error (E) or exit code (X)</dd>
 *     <dt>unsigned short</dt>
 *     <dd>The length of the bytes to be sent.</dd>
 *     <dt>bytes</dt>
 *     <dd>The data</dd>
 * </dl>
 * The exit code chunk will always be the last chunk and there will only ever be one. There can be multiple chunks of
 * standard output or standard error and they can be interleaved.
 */
public class GitHookOutputWriter {

    static final int SOCKET_OUTPUT_MAX = 65535;

    private static final char TYPE_EXIT_CODE = 'X';
    private static final char TYPE_STDERR = 'E';
    private static final char TYPE_STDOUT = 'O';

    private final WritableByteChannel channel;

    private CacheResult cacheResult;
    private String filter;
    private HostingOperationType operationType;

    public GitHookOutputWriter(@Nonnull WritableByteChannel channel) {
        this.channel = requireNonNull(channel, "channel");
    }

    /**
     * Retrieves an {@link OutputStream} wrapping the underlying channel.
     * <p>
     * <b>Warning</b>: It is <i>required</i> that callers only use this stream for data that is already in the
     * required format; data written to the returned stream will <i>not</i> be formatted. Access to the stream
     * is provided to support integrating with the Content Cache, which requires access to a stream in order to
     * pump in-progress cache hits, where the cache will contain <i>pre-formatted</i> output.
     *
     * @return a stream which can be used to write <i>structured data</i> to the underlying channel
     */
    @Nonnull
    public OutputStream asOutputStream() {
        return Channels.newOutputStream(channel);
    }

    /**
     * Sets the cache result that should be applied to the surrounding request when hook handling completes.
     * It will be applied to the request even if hook handling ultimately fails.
     *
     * @param cacheResult the cache result to set
     */
    public void setCacheResult(@Nonnull CacheResult cacheResult) {
        this.cacheResult = requireNonNull(cacheResult, "cacheResult");
    }

    /**
     * Sets the object filter that was used.
     *
     * @param filter the filter used, which may be {@code null}
     */
    public void setFilter(@Nullable String filter) {
        this.filter = filter;
    }

    /**
     * Sets the operation type that should be applied to the surrounding request when hook handling completes.
     * It will be applied to the request even if hook handling ultimately fails.
     *
     * @param operationType the operation type to set
     */
    public void setOperationType(@Nonnull HostingOperationType operationType) {
        this.operationType = requireNonNull(operationType, "operationType");
    }

    public void writeExitCode(int exitCode) throws IOException {
        writeAsciiChar(TYPE_EXIT_CODE);
        writeShort(exitCode);
    }

    public void writeStderr(@Nonnull ByteBuffer buffer) throws IOException {
        writeInChunks(TYPE_STDERR, requireNonNull(buffer, "buffer"));
    }

    public void writeStderr(@Nonnull String value) throws IOException {
        writeInChunks(TYPE_STDERR, requireNonNull(value, "value"));
    }

    public void writeStdout(@Nonnull ByteBuffer buffer) throws IOException {
        writeInChunks(TYPE_STDOUT, requireNonNull(buffer, "buffer"));
    }

    public void writeStdout(@Nonnull String value) throws IOException {
        writeInChunks(TYPE_STDOUT, requireNonNull(value, "value"));
    }

    @Nullable
    CacheResult getCacheResult() {
        return cacheResult;
    }

    @Nullable
    String getFilter() {
        return filter;
    }

    @Nullable
    HostingOperationType getOperationType() {
        return operationType;
    }

    private void writeAsciiChar(char c) throws IOException {
        checkArgument((c >= 0x0001) && (c <= 0x007F), "Only ASCII characters are supported");

        ByteBuffer buffer = ByteBuffer.allocate(1);
        buffer.put((byte) c);
        buffer.flip();
        channel.write(buffer);
    }

    /**
     * Writes the provided <i>binary data</i> in chunks to the channel.
     * <p>
     * <b>Warning</b>: If the binary data being written is encoded UTF-8, the chunking applied here is <i>not</i>
     * multibyte-aware and may split buffers in the middle of a character. Callers should either ensure they only
     * pass in buffers that will fit in a single packet, or use {@link #writeInChunks(char, String)} instead.
     *
     * @param type the chunk type being written
     * @param data the data to output as one or more chunks
     * @throws IOException if the data cannot be written
     */
    private void writeInChunks(char type, ByteBuffer data) throws IOException {
        do {
            //Write the channel
            writeAsciiChar(type);

            //Write the chunk size
            int size = Math.min(SOCKET_OUTPUT_MAX, data.remaining());
            writeShort(size);

            //Write the chunk's data
            ByteBuffer buffer = (ByteBuffer) data.slice().limit(size);
            channel.write(buffer);

            data.position(data.position() + buffer.limit());
        } while (data.hasRemaining());
    }

    /**
     * Uses a {@code CharsetEncoder} to perform multibyte-aware encoding of the provided data and write it
     * in chunks to the channel. {@code CharsetEncoder} will never write partial bytes for what should be a
     * single character to the output {@code ByteBuffer}, ensuring chunks contain complete UTF-8 codepoints
     * for all characters.
     *
     * @param type the chunk type being written
     * @param data the data to output as one or more UTF-8-encoded chunks
     * @throws IOException if the data cannot be encoded, or the encoded data cannot be written
     */
    private void writeInChunks(char type, String data) throws IOException {
        //Using REPLACE here matches what String.getBytes(StandardCharsets.UTF_8) would do
        CharsetEncoder utf8Encoder = StandardCharsets.UTF_8.newEncoder()
                .onMalformedInput(CodingErrorAction.REPLACE)
                .onUnmappableCharacter(CodingErrorAction.REPLACE);

        CharBuffer chars = CharBuffer.wrap(data);
        ByteBuffer bytes = ByteBuffer.allocate(Math.min(SOCKET_OUTPUT_MAX, chars.length() * 4));

        //Note that this WritableByteChannel is not buffered, so it's safe to intermix calls directly to
        //the DataOutputStream with calls to the channel. This just simplifies writing ByteBuffers
        CoderResult result;
        do {
            //Encode a chunk of characters. This will automatically stop when a full character cannot
            //fit in the remaining bytes; no need to manually try to back off the buffer
            result = utf8Encoder.encode(chars, bytes, true);
            bytes.flip();

            //Write the channel, then the packet size and the encoded bytes
            writeAsciiChar(type);
            writeShort(bytes.remaining());
            channel.write(bytes);

            //Clear the encoded bytes to prepare for encoding the next chunk
            bytes.clear();
        } while (result.isOverflow());
    }

    private void writeShort(int value) throws IOException {
        ByteBuffer buffer = ByteBuffer.allocate(2);
        buffer.putShort((short) value);
        buffer.flip();
        channel.write(buffer);
    }
}
