package com.atlassian.bitbucket.mesh.util;

import com.atlassian.bitbucket.mesh.git.builder.catfile.BatchHeader;

import javax.annotation.Nonnull;
import java.nio.ByteBuffer;

import static java.nio.charset.StandardCharsets.UTF_8;

/**
 * A {@link ByteBuffer} scanner which searches for a specified terminator byte, using a buffer with a specified
 * maximum size.
 */
public class ByteSegmentScanner {

    private final int maxSegmentSize;
    private final int terminator;

    /**
     * A convenience constructor to create a new scanner that uses line feeds (LF) as terminators with a max segment
     * size of {@link BatchHeader#MAX_SIZE_IN_BYTES}, for use detecting {@link BatchHeader}s.
     * <p>
     * This constructor is equivalent to:
     * <code>
     *     new ByteSegmentScanner((byte) '\n', BatchHeader.MAX_SIZE_IN_BYTES);
     * </code>
     */
    public ByteSegmentScanner() {
        this((byte) '\n', BatchHeader.MAX_SIZE_IN_BYTES);
    }

    /**
     * @param terminator    the terminator to use when looking for segments
     * @param maxSegmentSize maximum segment size
     */
    public ByteSegmentScanner(byte terminator, int maxSegmentSize) {
        this.maxSegmentSize = maxSegmentSize;
        this.terminator = terminator;
    }

    /**
     * @param output        the output to scan for a segment
     * @param closed        {@code true} if this will be the final output scanned; otherwise, {@code false} if
     *                      additional output may be provided
     * @return the scanned segment, or {@code null} if the terminator was not found {@code closed} is {@code false}
     */
    public ByteBuffer scan(@Nonnull ByteBuffer output, boolean closed) {
        if (output.remaining() == 0) {
            return null;
        }

        //This loop looks for the terminator past the maxSegmentSize because for missing objects where a long
        //object identifier was used (e.g. HEAD:some-long-path), the "missing" response will repeat that object
        //identifier.
        int limit = output.limit();
        for (int i = output.position(); i < limit; ++i) {
            if (output.get(i) == terminator) {
                ByteBuffer segment = output.slice();
                segment.limit(i - output.position());

                //Advance the output to the end of the segment, plus 1 byte for the terminator
                output.position(i + 1);

                return segment;
            }
        }

        int remaining = output.remaining();
        if (closed && remaining <= maxSegmentSize) {
            //If we didn't find a terminator, but this is the final output and the output is within the
            //specified limit, consume the output and return its contents as the segment
            ByteBuffer segment = output.slice();
            output.position(output.position() + output.remaining());

            return segment;
        }

        if (remaining > maxSegmentSize) {
            //If we get here, no terminator was found in the remainder of output, which is unexpected. Include the
            //full message in the exception message so we have useful information in the logs.
            output.mark();
            String actualData = UTF_8.decode(output).toString();
            output.reset();

            throw new IllegalStateException("Maximum segment size of " + maxSegmentSize + " exceeded. Data was '" +
                    actualData + "'");
        }

        //Otherwise, if we make it here, it means the output contained fewer bytes than the specified
        //limit and was not closed, and the terminator was not found.
        //
        //Note: The output's position has not been advanced.
        return null;
    }
}