package com.atlassian.stash.internal.hook;

import java.io.*;
import java.nio.charset.StandardCharsets;

import static com.google.common.base.Preconditions.checkArgument;

/**
 * An abstraction of the Stash <em>SCM hook output protocol</em>.
 * <p>
 * Output data over the socket is expected as a series of chunks in the following format:
 * <code>[ASCII char][unsigned short][UTF8 string]</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 either standard output (O), standard error (E) or the exit code (X)</dd>
 *     <dt>unsigned short</dt>
 *     <dd>The length of the UTF8 string to be sent.</dd>
 *     <dt>UTF8 string</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 SocketTransferOutput implements Flushable, Closeable {

    static final int SOCKET_OUTPUT_MAX = 65535;

    private final DataOutputStream out;

    public SocketTransferOutput(OutputStream out) {
        this.out = new DataOutputStream(out);
    }

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

    public void writeStdOut(String chunk) throws IOException {
        writeInChunks('O', chunk);
    }

    public void writeStdErr(String chunk) throws IOException {
        writeInChunks('E', chunk);
    }

    public void writeExitCode(int exitCode) throws IOException {
        writeAsciiChar('X');
        writeUnsignedShort(exitCode);
    }

    private void writeInChunks(char channel, String data) throws IOException {
        byte[] bytes = data.getBytes(StandardCharsets.UTF_8);
        int start = 0;
        while (start < bytes.length) {
            // Chunks can be at most SOCKET_OUTPUT_MAX bytes long, because the length is written as 2 bytes (short).
            int end = rewindToLastBoundary(bytes, start + SOCKET_OUTPUT_MAX);
            int len = end - start;
            writeAsciiChar(channel);
            // Note that we cannot use writeUTF here because it writes "modified UTF-8", which is different from UTF-8
            out.writeShort(len);
            out.write(bytes, start, len);

            start = end;
        }
    }

    /**
     * UTF-8 is variable width and uses the following system:
     * <ul>
     * <li>Single-byte characters are a 0b0* byte</li>
     * <li>Multi-byte characters start with a 0b11* byte and continue with 0b10* bytes</li>
     * </ul>
     * To not split a multi-byte character into two, all we have to do is rewind to the first non-continuation byte.
     */
    private static int rewindToLastBoundary(byte[] bytes, int end) {
        if (end >= bytes.length) {
            return bytes.length;
        }

        while ((bytes[end] & 0b11000000) == 0b10000000) {
            end--;
        }
        return end;
    }

    private void writeUnsignedShort(int positiveInt) throws IOException {
        checkArgument(positiveInt >= 0, "positiveInt must be positive");
        out.write((positiveInt >>> 8) & 0xFF);
        out.write(positiveInt & 0xFF);
    }

    @Override
    public void flush() throws IOException {
        out.flush();
    }

    @Override
    public void close() throws IOException {
        out.close();
    }
}
