package com.atlassian.bitbucket.mesh.process.handler;

import javax.annotation.Nonnull;
import javax.annotation.concurrent.NotThreadSafe;
import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.charset.Charset;
import java.nio.charset.CharsetDecoder;
import java.nio.charset.CoderResult;
import java.nio.charset.CodingErrorAction;

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

/**
 * A utility class for taking in {@link ByteBuffer}s, decoding them into a specified {@link Charset}, and then
 * dispatching them to a provided {@link LineSink sink}.
 */
@NotThreadSafe
public class LineDecoder {

    private static final char NULL_CHAR = (char) 0x00;

    private final ReadAheadCharSource charSource;
    private final CharsetDecoder decoder;
    private final Mode mode;

    private StringBuilder builder;
    private int cutoffLength;
    private int maxLineLength;

    public LineDecoder(@Nonnull Charset charset) {
        this(charset, Mode.CRLF_OR_LF);
    }

    public LineDecoder(@Nonnull Charset charset, @Nonnull Mode mode) {
        this(charset, mode, Integer.MAX_VALUE);
    }

    public LineDecoder(@Nonnull Charset charset, @Nonnull Mode mode, int maxLineLength) {
        this.mode = requireNonNull(mode, "mode");

        builder = new StringBuilder();
        decoder = requireNonNull(charset, "charset").newDecoder()
                //Using REPLACE for these actions matches what the StreamDecoder created by an
                //InputStreamReader does, which keeps BIO and NIO handling consistent
                .onMalformedInput(CodingErrorAction.REPLACE)
                .onUnmappableCharacter(CodingErrorAction.REPLACE);
        charSource = new ReadAheadCharSource();

        setMaxLineLength(maxLineLength);
    }

    /**
     * Decodes the provided {@link ByteBuffer}, then assembles the resulting characters into lines and dispatches
     * them to the provided {@link LineSink sink}.
     * <p>
     * If the sink returns {@code false}, any further characters (or undecoded bytes) will be ignored and the
     * method will return {@code true} to indicate the sink is done. Otherwise, the method will return {@code closed}.
     * <p>
     * If the provided buffer contains an incomplete line (i.e. a line with no EOL marker), what happens depends
     * on the {@code closed} flag:
     * <ul>
     *     <li>{@code true}: It is assumed the output is missing a trailing EOL (not uncommon) and the EOS (end of
     *     stream) is treated as EOL</li>
     *     <li>{@code false}: The rest of the line may be a subsequent buffer, so the partial line will be
     *     <i>buffered internally</i> pending a subsequent call to {@code forEach}.</li>
     * </ul>
     * The provided {@link ByteBuffer} is always <i>fully consumed</i> <i>unless</i> the {@link LineSink sink}
     * returns {@code false}.
     *
     * @param bytes  bytes to decode and dispatch to the provided sink
     * @param closed {@code true} to indicate there will be no more bytes after the provided buffer
     * @param sink   receives each decoded line and a boolean to indicate if the line was truncated, and may return
     *               {@code false} to indicate no more lines should be provided
     * @return {@code true} to indicate decoding is <i>complete</i> (meaning either {@code closed} was {@code true},
     *         or that the {@link LineSink sink} returned {@code false} to indicate it does not want more lines;
     *         otherwise, {@code false} to indicate decoding should continue
     */
    public boolean forEach(@Nonnull ByteBuffer bytes, boolean closed, @Nonnull LineSink sink) {
        requireNonNull(sink, "sink");

        //Prepare a CharBuffer to hold up to 5000 characters, or fewer if the ByteBuffer is shorter.
        //It's possible each character will consume more than one byte, but this sizing handles the
        //worst-case scenario where each character is a single byte.
        CharBuffer chars = CharBuffer.allocate(Math.min(5000, requireNonNull(bytes, "bytes").remaining()));

        CoderResult result;
        do {
            //Decode the next set of characters and flip the buffer to prepare for reading
            result = decoder.decode(bytes, chars, closed);
            chars.flip();

            charSource.fill(chars);
            int c;
            while ((c = charSource.readNext()) != -1) {
                boolean isEol = false;
                if (mode == Mode.LF) {
                    isEol = c == '\n';
                } else if (mode == Mode.NULL) {
                    isEol = c == NULL_CHAR;
                } else {
                    if (c == '\n') {
                        isEol = true;
                    } else if (c == '\r') {
                        //For these modes, CR might be interpreted as a new line depending on the following character.
                        int c2 = charSource.readNext();
                        if (c2 == -1) {
                            //If the block ends here, we push the character we just read back and break the loop, either
                            //causing the next block to be read, or the remainder of the block being fed to the line sink.
                            charSource.pushBack(c);
                            break;
                        } else if (c2 == '\n') {
                            //CR followed by LF definitely means we have a new line, so create one
                            isEol = true;
                        } else if (mode == Mode.CRLF_CR_OR_LF) {
                            //CR not followed by LF is only a new line in this mode. Create a new line and push back the
                            //character that would otherwise be missing.
                            isEol = true;
                            charSource.pushBack(c2);
                        } else {
                            //CR not followed by LF. Push back the character so it can be read again.
                            charSource.pushBack(c2);
                        }
                    }
                }
                if (isEol) {
                    if (offerLine(sink, builder)) {
                        builder = new StringBuilder(); //Allocate a new builder in case the buffer was large
                    } else {
                        return true;
                    }
                } else if (builder.length() < cutoffLength) {
                    builder.append((char) c);
                }
            }

            chars.clear();
        } while (result.isOverflow()); //If the ByteBuffer is too long, we may need to decode multiple times

        //This handles the case where the output doesn't have a trailing EOL marker to ensure the final line
        //of output isn't dropped. If the stream is closed and the builder isn't empty, treat it as a line
        if (closed && builder.length() != 0) {
            //The sink's response here doesn't matter; whether they would accept more lines or not, there
            //are no more lines to offer (because the buffer is empty and the stream is closed)
            offerLine(sink, builder);
        }

        return closed;
    }

    public int getMaxLineLength() {
        return maxLineLength;
    }

    public void setMaxLineLength(int maxLineLength) {
        this.maxLineLength = maxLineLength < 1 ? Integer.MAX_VALUE : maxLineLength;

        if (this.maxLineLength > Integer.MAX_VALUE - 2) { //Check for integer overflow
            cutoffLength = Integer.MAX_VALUE;
        } else {
            //+2 is CR (to ensure CRLF handling works correctly) + 1 character to detect truncation
            cutoffLength = this.maxLineLength + 2;
        }
    }

    /**
     * @return {@code true} to indicate that more input is desired, {@code false} to signal to stop.
     */
    private boolean offerLine(LineSink sink, StringBuilder line) {
        boolean truncated = line.length() > maxLineLength;

        return sink.accept(truncated ? line.substring(0, maxLineLength) : line.toString(), truncated);
    }

    public enum Mode {

        CRLF_OR_LF, // Use \n or \r\n as line separators
        CRLF_CR_OR_LF, // Use \r, \n or \r\n as line separators
        LF, // Use \n characters as line separators
        NULL; // Use NUL characters as line separators
    }

    private static class ReadAheadCharSource {

        CharBuffer currentBuffer;
        int readAhead;

        public ReadAheadCharSource() {
            readAhead = -1;
        }

        void fill(CharBuffer buffer) {
            currentBuffer = buffer;
        }

        void pushBack(int readAhead) {
            checkState(this.readAhead == -1, "read ahead buffer is not empty");

            this.readAhead = readAhead;
        }

        int readNext() {
            if (readAhead != -1) {
                int readAhead = this.readAhead;
                this.readAhead = -1;
                return readAhead;
            }
            if (!currentBuffer.hasRemaining()) {
                return -1;
            }
            return currentBuffer.get();
        }
    }
}
