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

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import javax.annotation.concurrent.NotThreadSafe;
import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.charset.Charset;
import java.nio.charset.CharsetEncoder;
import java.util.Iterator;
import java.util.stream.Stream;

import static java.util.Objects.requireNonNull;

/**
 * A utility class that takes a list of lines, encodes them using the provided {@link Charset} and emits the encoded
 * bytes to a provided {@link ByteBuffer}. This class should be used to encode lines of input data and emit them to a
 * process's {@code stdin}.
 */
@NotThreadSafe
public class LineEncoder {

    private final Charset charset;
    private final Iterator<String> lineIterator;
    private final Mode mode;

    private CharsetEncoder encoder;
    private CharBuffer nextFragment;

    public LineEncoder(@Nonnull Iterable<String> lines, @Nonnull Charset charset) {
        this(requireNonNull(lines, "lines").iterator(), charset);
    }

    public LineEncoder(@Nonnull Stream<String> lines, @Nonnull Charset charset) {
        this(requireNonNull(lines, "lines").iterator(), charset);
    }

    public LineEncoder(@Nonnull Iterator<String> lineIterator, @Nonnull Charset charset) {
        this(lineIterator, charset, Mode.LF);
    }

    public LineEncoder(@Nonnull Iterator<String> lineIterator, @Nonnull Charset charset, @Nullable Mode mode) {
        this.charset = requireNonNull(charset, "charset");
        this.lineIterator = requireNonNull(lineIterator, "lineIterator");
        this.mode = mode == null ? Mode.LF : mode;
    }

    /**
     * Encodes and writes the next portion of the input {@code lines} to the provided {@code buffer}.
     *
     * @param buffer a buffer to write the encoded bytes to
     * @return {@code true} to indicate more data is available, which did not fit in the buffer; otherwise,
     *         {@code false} to indicate there's no more data available
     */
    public boolean encodeTo(@Nonnull ByteBuffer buffer) {
        if (encoder == null) {
            //Prepare the encoder
            encoder = charset.newEncoder();
        }
        if (nextFragment != null) {
            if (encoder.encode(nextFragment, buffer, !lineIterator.hasNext()).isOverflow()) {
                //If the line won't fit in the buffer, retain the nextFragment, flip the buffer and
                //return. The next callback will continue writing the line
                buffer.flip();

                return true;
            }
            //Otherwise, the line fit into the buffer and we can move on
            nextFragment = null;
        }

        while (lineIterator.hasNext()) {
            CharBuffer line = CharBuffer.wrap(lineIterator.next() + mode.getSeparator());
            if (encoder.encode(line, buffer, !lineIterator.hasNext()).isOverflow()) {
                buffer.flip();
                nextFragment = line;

                return true;
            }
        }

        buffer.flip();
        return false;
    }

    public boolean hasMore() {
        return nextFragment != null || lineIterator.hasNext();
    }

    public enum Mode {

        LF("\n"), // Use \n characters as line separators
        NULL("\u0000"); // Use NUL characters as line separators

        private final String separator;

        Mode(String separator) {
            this.separator = separator;
        }

        @Nonnull
        public String getSeparator() {
            return separator;
        }
    }
}
