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

import com.atlassian.bitbucket.mesh.rpc.ZeroCopyUtils;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamChunk;
import com.atlassian.bitbucket.mesh.util.SatiableConsumer;
import com.google.protobuf.ByteString;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.zip.GZIPInputStream;

import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.Objects.requireNonNull;

/**
 * Stream that consumes a stream of {@link RpcStreamChunk}, identifies git packet lines and the PACK section, filters
 * them through {@link #addFilter(GitReceivePackFilter) registered filters} and emits the result to a delegate
 * consumer.
 * <p>
 * Registered filters are invoked in the order in which they have been registered. The filter's output is used as
 * input to the next filter, and eventually as input to the delegate consumer.
 */
public class FilteringGitReceivePackStream {

    private static final byte[] PACK_HEADER = "PACK".getBytes(UTF_8);
    private static final Logger log = LoggerFactory.getLogger(FilteringGitReceivePackStream.class);
    private final SatiableConsumer<RpcStreamChunk> delegate;
    private final List<GitReceivePackFilter> filters;
    private final byte[] packetPrefix;

    //The number of bytes to the start of the next packet line from the current read position in the current ByteBuffer
    private int bytesToNextPacket;
    private boolean delegateIsDone;
    private boolean filterPack;
    private boolean filterPacketLines;
    private boolean inPack;
    private boolean invalid;
    private int packetPrefixIndex;
    private ByteString partialPacketLine;

    public FilteringGitReceivePackStream(SatiableConsumer<RpcStreamChunk> delegate) {
        this.delegate = delegate;

        filters = new ArrayList<>();
        packetPrefix = new byte[4];
    }

    public void addFilter(@Nonnull GitReceivePackFilter filter) {
        filters.add(requireNonNull(filter, "filter"));
        filterPack |= filter.wantsPack();
        filterPacketLines |= filter.wantsPacketLines();
    }

    public void removeFilter(@Nonnull GitReceivePackFilter filter) {
        filters.remove(requireNonNull(filter, "filter"));
        filterPack = filters.stream().anyMatch(GitReceivePackFilter::wantsPack);
        filterPacketLines = filters.stream().anyMatch(GitReceivePackFilter::wantsPacketLines);
    }

    public void write(@Nonnull RpcStreamChunk chunk) {
        if (delegateIsDone) {
            return;
        }
        if (invalid || !(filterPack || filterPacketLines)) {
            //No point in filtering if it's not a git packet line stream or none of the filters are interested anymore
            emitChunk(chunk);
        } else if (inPack) {
            //Once the PACK section starts, the remainder of the stream consists of PACK chunks. The chunk does not
            //have to be analyzed and can be handed over to the filters, and the result written to the delegate.
            emitChunk(filterPack(chunk));
        } else {
            //We're still in the git packet-line section of the receive-pack protocol. Analyze the chunk to identify
            //git packet lines, and possibly the start of the PACK section. All identified sub-chunks are
            //filtered, and the filter result written to the delegate.
            doFilter(chunk);

            if (chunk.getClosed()) {
                emitLastChunk();
            }
        }
    }

    private static int getPacketSize(byte[] packetPrefix) {
        //The flush (0000) and delim (0001) packets are special because they don't encode their true length
        //in bytes. Therefore we apply a minimum size of 4.
        return Math.max(4, Integer.parseInt(new String(packetPrefix, UTF_8), 16));
    }

    private static boolean isGzipped(byte[] bytes) {
        int first = bytes[0] & 0xFF;  //The & is to convert the signed byte to unsigned
        int second = bytes[1] & 0XFF; //to match the int result of InputStream.read()
        int magic = (second << 8) | first; //| the two together to make an unsigned short

        return magic == GZIPInputStream.GZIP_MAGIC;
    }

    private static boolean isPackHeader(byte[] packetPrefix) {
        return Arrays.equals(PACK_HEADER, packetPrefix);
    }

    private void bufferPacketLine(RpcStreamChunk chunk, int startIndex, int endIndex) {
        if (startIndex < endIndex) {
            //The provided RpcStreamChunk may be backed by direct memory that is returned to the Netty buffer pool
            //at some point after this method returns. Create a deep copy because we're buffering between different
            //invocations of write(RpcStreamChunk).
            ByteString data = ZeroCopyUtils.deepCopy(partialData(chunk, startIndex, endIndex));
            partialPacketLine = partialPacketLine == null ? data : partialPacketLine.concat(data);
        }
    }

    private void doFilter(RpcStreamChunk chunk) {
        ByteString data = chunk.getData();
        int chunkSize = data.size();
        //If packetPrefixIndex < 0, there's data in packetPrefix that we haven't emitted yet
        int emittedIndex = -packetPrefixIndex;
        int indexInChunk = emittedIndex;
        for (ByteBuffer buffer : data.asReadOnlyByteBufferList()) {
            int bufferSize = buffer.remaining();
            if (bytesToNextPacket >= bufferSize) {
                //The next packet line starts beyond buffer so move on to the next buffer
                bytesToNextPacket -= bufferSize;
                indexInChunk += bufferSize;
                continue;
            }

            while (!delegateIsDone && buffer.hasRemaining()) {
                if (bytesToNextPacket > 0) {
                    //Advance position to the start of the next packet or to the end of the buffer,
                    //whichever is first
                    int skippedBytes = Math.min(bytesToNextPacket, buffer.remaining());
                    buffer.position(buffer.position() + skippedBytes);
                    bytesToNextPacket -= skippedBytes;
                    indexInChunk += skippedBytes;
                }

                if (buffer.hasRemaining()) {
                    //We're positioned at the packet line prefix; emit the packet line leading up to the prefix.
                    //If no filters are interested in packet lines, this will no-op and not emit the data.
                    emittedIndex += filterAndEmitPacketLine(chunk, emittedIndex, indexInChunk);
                    //If bytesToNextPacket >= 0, we're positioned at the start of the packet line prefix.
                    //If bytesToNextPacket < 0, we've already read a partial prefix into packetPrefix and
                    //are positioned at the start of the remainder of the prefix.
                    readPrefix(buffer);
                    if (packetPrefixIndex == 4) {
                        //We've read the full prefix
                        if (isPackHeader(packetPrefix)) {
                            //Emit the data leading up to the pack header. If any of the filters are processing
                            //packet lines, this will be empty because of the filterAndEmitPacketLine call above.
                            //If none of the filters are processing packet lines, there may be one or more packet lines
                            //that have not yet been emitted. Emit them in a single chunk.
                            if (indexInChunk != emittedIndex) {
                                emitChunk(partialChunk(chunk, emittedIndex, indexInChunk));
                            }
                            inPack = true;
                            //Offer the start of the PACK section up to the end of the chunk to the filters
                            //and emit the result.
                            emitChunk(filterPack(partialChunk(chunk, indexInChunk, chunkSize)));
                            packetPrefixIndex = 0;
                            return;
                        }
                        try {
                            //We've found the next packet line and have read its prefix into packetPrefix
                            int currentPacketSize = getPacketSize(packetPrefix);
                            bytesToNextPacket = currentPacketSize - 4; //-4 because we've already read the 4 byte prefix
                            indexInChunk += 4; //+4 because we've read the prefix
                            packetPrefixIndex = 0;
                        } catch (NumberFormatException e) {
                            //The remainder of the chunk is invalid, forward it to the delegate unfiltered
                            RpcStreamChunk invalidRemainder = partialChunk(chunk, emittedIndex, chunkSize);
                            markInvalid();
                            emitChunk(invalidRemainder);

                            if (emittedIndex == 0 && isGzipped(packetPrefix)) {
                                log.debug("Disabling receive-pack filtering because the data stream is compressed");
                            } else {
                                log.warn("Unexpected packet line prefix '{}'. This does not look like a packet line: '{}'",
                                        new String(packetPrefix, UTF_8), invalidRemainder.getData().toStringUtf8());
                            }
                            return;
                        }
                    } else {
                        //If we get here, we've reached the end of buffer and have ended with a partially read
                        //packet line prefix. That prefix runs into the next buffer, or even the next chunk.
                        //Update bytesToNextPacket to be a _negative_ value to remember that part of packet lies
                        //in the (soon to be) previous buffer and/or chunk.
                        bytesToNextPacket = -packetPrefixIndex;
                    }
                }
            }
        }

        if (!delegateIsDone) {
            //If none of the filters are filtering packet lines, we can emit the part of chunk that we know
            //consists of packet lines.
            emittedIndex += emitUnfilteredPacketLines(chunk, emittedIndex, indexInChunk);
            if (bytesToNextPacket == 0 && indexInChunk > emittedIndex) {
                //The chunk ended with a 4 byte packet line (FLUSH or DELIM). Emit the packet
                emittedIndex += filterAndEmitPacketLine(chunk, emittedIndex, indexInChunk);
            }
            if (emittedIndex < indexInChunk) {
                bufferPacketLine(chunk, emittedIndex, indexInChunk);
            }
        }
    }

    private void emitChunk(RpcStreamChunk chunk) {
        if (!delegateIsDone && chunk != null) {
            delegateIsDone = delegate.accept(chunk);
        }
    }

    private void emitLastChunk() {
        if (partialPacketLine != null) {
            emitChunk(RpcStreamChunk.newBuilder()
                    .setData(partialPacketLine)
                    .setClosed(true)
                    .build());
            partialPacketLine = null;
        } else if (packetPrefixIndex > 0) {
            emitChunk(RpcStreamChunk.newBuilder()
                    .setData(ByteString.copyFrom(packetPrefix, 0, packetPrefixIndex))
                    .setClosed(true)
                    .build());
        }
    }

    /**
     * @return the number of bytes of the provided chunk that has been emitted
     */
    private int emitUnfilteredPacketLines(RpcStreamChunk chunk, int startIndex, int endIndex) {
        if (filterPacketLines || delegateIsDone || startIndex == endIndex) {
            return 0;
        }
        //We're still in the packet line section, but have not emitted anything to the filters
        //because no filter is interested in packet lines. Any completely processed packet lines in chunk
        //can be written to delegate.
        int chunkSize = chunk.getData().size();
        if (startIndex == 0 && endIndex == chunkSize) {
            //We can write the chunk unchanged to the delegate
            emitChunk(chunk);
            return chunkSize;
        }

        //Chunk starts (negative startIndex) and/or ends (endIndex < chunkSize) with a partial packet prefix.
        //Include a previously stored partial prefix in the outgoing chunk.
        ByteString data = chunk.getData().substring(Math.max(0, startIndex), endIndex);
        if (partialPacketLine != null) {
            data = partialPacketLine.concat(data);
            partialPacketLine = null;
        }
        if (endIndex < chunkSize) {
            //Chunk ends with a partial packet line. Store it in partialPacketLine so it can be emitted
            //when we have the full prefix and can determine whether we've reached the end of the packet lines
            partialPacketLine = chunk.getData().substring(endIndex, chunkSize);
        }
        emitChunk(RpcStreamChunk.newBuilder()
                .setClosed(chunk.getClosed() && endIndex == chunkSize)
                .setData(data)
                .build());
        return data.size();
    }

    /**
     * @return the number of bytes from the provided chunk that has been filtered and emitted
     */
    private int filterAndEmitPacketLine(RpcStreamChunk chunk, int startIndex, int endIndex) {
        if (filterPacketLines && (startIndex != endIndex || partialPacketLine != null)) {
            ByteString packetLine = partialData(chunk, startIndex, endIndex);
            RpcStreamChunk packetLineChunk = RpcStreamChunk.newBuilder()
                    .setClosed(chunk.getClosed() && endIndex == chunk.getData().size())
                    .setData(partialPacketLine == null ? packetLine : partialPacketLine.concat(packetLine))
                    .build();
            partialPacketLine = null;

            emitChunk(filterPacketLines ? filterPacketLine(packetLineChunk) : packetLineChunk);
            return packetLine.size();
        }
        //If no filters are interested in packet lines, don't bother constructing a sub-chunk and emitting it.
        //Instead, we'll emit all packet lines in the chunk as a combined chunk. See emitUnfilteredPacketLines.
        return 0;
    }

    private RpcStreamChunk filterPack(RpcStreamChunk chunk) {
        if (!delegateIsDone) {
            boolean continueFiltering = false;
            for (GitReceivePackFilter filter : filters) {
                if (filter.wantsPack()) {
                    try {
                        chunk = filter.onPack(chunk);
                        continueFiltering |= filter.wantsPack();
                    } catch (RuntimeException e) {
                        log.warn("Filter {} failed to transform pack chunk", filter, e);
                    }
                }
            }
            filterPack = continueFiltering;
        }
        return chunk;
    }

    private RpcStreamChunk filterPacketLine(RpcStreamChunk chunk) {
        if (!delegateIsDone) {
            boolean continueFiltering = false;
            for (GitReceivePackFilter filter : filters) {
                if (filter.wantsPacketLines()) {
                    try {
                        chunk = filter.onPacketLine(chunk);
                        continueFiltering |= filter.wantsPacketLines();
                    } catch (RuntimeException e) {
                        log.warn("Filter {} failed to transform pack line", filter, e);
                    }
                }
            }
            filterPacketLines = continueFiltering;
        }
        return chunk;
    }

    private void markInvalid() {
        invalid = true;
        for (GitReceivePackFilter filter : filters) {
            try {
                filter.onInvalid();
            } catch (RuntimeException e) {
                log.warn("Filter {} failed when notified that the stream is invalid", filter, e);
            }
        }
    }

    private RpcStreamChunk partialChunk(RpcStreamChunk chunk, int startIndex, int endIndex) {
        if (startIndex == 0) {
            return chunk;
        }
        return RpcStreamChunk.newBuilder()
                .setClosed(chunk.getClosed() && endIndex == chunk.getData().size())
                .setData(partialData(chunk, startIndex, endIndex))
                .build();
    }

    private ByteString partialData(RpcStreamChunk chunk, int startIndex, int endIndex) {
        if (startIndex < 0 && startIndex >= -4) {
            //A negative startIndex means that the packet line header bytes were split between this chunk and
            //the previous chunk, and that the missing 1..4 bytes can be found in the packetPrefix buffer.
            ByteString missingHeader = ByteString.copyFrom(packetPrefix, 0, -startIndex);
            return missingHeader.concat(chunk.getData().substring(0, endIndex));
        }
        return chunk.getData().substring(startIndex, endIndex);
    }

    private void readPrefix(ByteBuffer buffer) {
        int read = Math.min(buffer.remaining(), 4 - packetPrefixIndex);
        buffer.get(packetPrefix, packetPrefixIndex, read);
        packetPrefixIndex += read;
    }
}
