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

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

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

import static com.google.common.base.Preconditions.checkArgument;
import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.Objects.requireNonNull;

/**
 * Segments a byte stream into a sequence of git packets. Users of this class can append bytes to the
 * {@code GitPacketStream}. When a complete git packet has been appended, the
 * {@link #addListener(GitPacketListener) listeners} of the stream get a
 * {@link GitPacketListener#onPacket(byte[]) callback}.
 */
public class GitPacketStream {

    private static final Logger log = LoggerFactory.getLogger(GitPacketStream.class);

    private static final int DEFAULT_BUFFER_SIZE = 4096;
    private static final byte[][] MARKER_PACKETS = {GitPackets.getFlush(), GitPackets.getDelim()};
    private static final int PACKET_PREFIX_SIZE = 4;
    private static final String PACK_HEADER = "PACK";

    private final boolean abortOnPackSection;
    private final int bufferSize;
    private final List<GitPacketListener> listeners;
    private final ByteArrayOutputStream partialPacket;

    private volatile boolean aborted;
    private volatile byte[] buffer;
    private boolean firstBuffer;
    private boolean inPackSection;
    private boolean inPrefix;
    private int nextPacketSize;

    public GitPacketStream() {
        this(true, DEFAULT_BUFFER_SIZE);
    }

    /**
     * @param abortOnPackSection whether to abort processing when the PACK section is detected
     * @param bufferSize the internal buffer size to use
     */
    public GitPacketStream(boolean abortOnPackSection, int bufferSize) {
        this.abortOnPackSection = abortOnPackSection;
        this.bufferSize = bufferSize;

        partialPacket = new ByteArrayOutputStream(1024);
        firstBuffer = true;
        inPrefix = true;
        listeners = new ArrayList<>();
        nextPacketSize = PACKET_PREFIX_SIZE;
    }

    public void abort() {
        aborted = true;
        partialPacket.reset();
        notifyOnAborted();
    }

    public void addListener(@Nonnull GitPacketListener listener) {
        listeners.add(requireNonNull(listener, "listener"));
    }

    public void append(@Nonnull byte[] bytes) {
        append(bytes, 0, bytes.length);
    }

    public void append(@Nonnull byte[] bytes, int off, int len) {
        requireNonNull(bytes, "bytes");
        checkArgument(off >= 0, "illegal offset");
        checkArgument(bytes.length >= off + len, "insufficient bytes");

        if (firstBuffer) {
            firstBuffer = false;

            if (len >= 2 && isGzipped(bytes)) {
                // The data is gzip-encoded. Currently we don't support analyzing gzip-encoded streams,
                // so abort to skip analysis (rather than allowing parsePacketSize to throw an exception because
                // it gets unexpected input)
                abort();
            }
        }
        if (aborted) {
            //Nothing left to do
            return;
        }

        int offset = off;
        int end = off + len;
        try {
            while (!aborted && offset < end && (nextPacketSize - partialPacket.size()) <= (end - offset)) {
                if (partialPacket.size() > 0) {
                    // buffer contains a partial git packet. Append to the buffer to complete the packet if possible.
                    int remainingPacketBytes = nextPacketSize - partialPacket.size();
                    partialPacket.write(bytes, offset, remainingPacketBytes);

                    // update offset to reflect that the 'bytes' array has been processed up to this point
                    offset += remainingPacketBytes;

                    if (inPrefix) {
                        readPacketPrefix(partialPacket.toByteArray(), 0);
                        if (aborted) {
                            return;
                        }
                        if (nextPacketSize <= 1) {
                            // FLUSH (0000) or DELIM (0001) packet encountered. Marker packets deviate from other git
                            // packets in that the packet prefix does not correspond with the actual length of the packet
                            // (0000 = 4 bytes).
                            readPacket(MARKER_PACKETS[nextPacketSize]);
                            nextPacketSize = PACKET_PREFIX_SIZE;
                            partialPacket.reset();
                            inPrefix = true;
                        }
                    } else {
                        readPacket(partialPacket.toByteArray());
                        partialPacket.reset();
                    }
                } else {
                    // no partial packet in buffer
                    if (inPrefix) {
                        readPacketPrefix(bytes, offset);
                        if (aborted) {
                            return;
                        }
                        if (nextPacketSize <= 1) {
                            // FLUSH (0000) or DELIM (0001) packet encountered. Marker packets deviate from other git
                            // packets in that the packet prefix does not correspond with the actual length of the packet
                            // (0000 = 4 bytes).
                            offset += readPacket(MARKER_PACKETS[nextPacketSize]);
                            nextPacketSize = PACKET_PREFIX_SIZE;
                            inPrefix = true;
                        }
                    } else {
                        byte[] packet = new byte[nextPacketSize];
                        System.arraycopy(bytes, offset, packet, 0, nextPacketSize);

                        offset += readPacket(packet);
                    }
                }
            }

            if (!aborted && offset < end) {
                partialPacket.write(bytes, offset, end - offset);
                offset = end;
            }
        } finally {
            if (offset != end) {
                // implies that an exception happened
                abort();
            }
        }
    }

    public void append(@Nonnull ByteBuffer byteBuffer) {
        if (!aborted) {
            if (buffer == null) {
                buffer = new byte[bufferSize];
            }

            while (!aborted && byteBuffer.hasRemaining()) {
                int length = Math.min(byteBuffer.remaining(), buffer.length);

                //Copy from byteBuffer into the byte[] buffer
                byteBuffer.get(buffer, 0, length);
                //Append from the byte[] buffer into the packet stream
                append(buffer, 0, length);
            }
        }
        //If processing was aborted, simply advance the position to the end of the buffer to 'consume' the whole buffer
        byteBuffer.position(byteBuffer.position() + byteBuffer.remaining());
    }

    public boolean isAborted() {
        return aborted;
    }

    public boolean isInPackSection() {
        return inPackSection;
    }

    /**
     * Inspects the provided {@code byte[]} to decide if it represents gzip-compressed data. If so, the bytes
     * can't be analyzed by the current implementation. This check should only be performed on the first array
     * of bytes read from the stream.
     *
     * @param bytes the bytes to check for gzip compression
     * @return {@code true} if the bytes start with the gzip magic constant; otherwise, {@code false} if the
     *         bytes appear to be uncompressed
     */
    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 void notifyListeners(byte[] packet) {
        for (GitPacketListener listener : listeners) {
            try {
                listener.onPacket(packet);
            } catch (Exception e) {
                // log the exception
                log.warn("Error while notifying {} of git packet", listener.getClass().getName(), e);
            }
        }
    }

    private void notifyOnAborted() {
        for (GitPacketListener listener : listeners) {
            try {
                listener.onAborted();
            } catch (Exception e) {
                // log the exception
                log.warn("Error while notifying {} of non-packet data", listener.getClass().getName(), e);
            }
        }
    }

    private int parsePacketSize(byte[] bytes, int offset) {
        String prefix = new String(bytes, offset, 4, UTF_8);
        if (!inPackSection && PACK_HEADER.equals(prefix)) {
            inPackSection = true;
            abort();
            return -1;
        }
        try {
            return Integer.parseInt(prefix, 16);
        } catch (NumberFormatException e) {
            throw new IllegalArgumentException("'" + prefix + "' is not a valid git packet prefix");
        }
    }

    private int readPacket(byte[] packet) {
        nextPacketSize = PACKET_PREFIX_SIZE;
        inPrefix = true;

        if (!inPackSection && packet.length > 8 &&
                packet[4] == 1 && packet[5] == 'P' && packet[6] == 'A' && packet[7] == 'C' && packet[8] == 'K') {
            inPackSection = true;
            if (abortOnPackSection) {
                abort();
            }
        }

        if (!aborted) {
            notifyListeners(packet);
        }
        return packet.length;
    }

    private void readPacketPrefix(byte[] bytes, int offset) {
        nextPacketSize = parsePacketSize(bytes, offset);
        inPrefix = false;
    }
}
