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

import com.atlassian.bitbucket.mesh.junit.LogListener;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamChunk;
import com.atlassian.bitbucket.mesh.util.SatiableConsumer;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.protobuf.ByteString;
import org.apache.commons.lang3.mutable.MutableInt;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.springframework.core.io.ClassPathResource;

import javax.annotation.Nonnull;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.UncheckedIOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static com.atlassian.bitbucket.mesh.test.MockitoUtils.returnFirst;
import static java.nio.charset.StandardCharsets.UTF_8;
import static org.junit.Assert.*;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;

public class FilteringGitReceivePackStreamTest {

    @Rule
    public final LogListener logListener = new LogListener(FilteringGitReceivePackStream.class);

    private CapturingRpcStream delegate;
    private CapturingFilter filter;
    private FilteringGitReceivePackStream stream;

    @Before
    public void setUp() {
        delegate = new CapturingRpcStream();
        filter = new CapturingFilter();
        stream = new FilteringGitReceivePackStream(delegate);
        stream.addFilter(filter);
    }

    @Test
    public void testFlushPacketHandled() {
        appendBytes("0006aa00000008abcdPACKtesting"); // 0006aa, 0000 and 0008abcd packets
        appendBytes(""); //empty chunk
        appendBytesAndClose("haha");

        assertPacketLines("0006aa", "0000", "0008abcd");
        assertPackEquals("PACKtestinghaha");
        assertFalse(filter.invalid);
        assertEquals("0006aa00000008abcdPACKtestinghaha", delegate.getData());
        assertTrue(delegate.closed);
    }

    @Test
    public void testFlushPacketAsLastPacketHandled() {
        appendBytesAndClose("0006aa0000"); // 0006aa and 0000 packets

        assertPacketLines("0006aa", "0000");
        assertFalse(filter.invalid);
        assertEquals("0006aa0000", delegate.getData());
        assertTrue(delegate.closed);
    }

    @Test
    public void testFilterCanModifyChunks() {
        stream = new FilteringGitReceivePackStream(delegate);
        stream.addFilter(new MutatingFilter("REPLACED", ImmutableMap.of(
                "0008abcd", "0009abcde",
                "0005a", "0006aa")));
        stream.addFilter(filter);

        appendBytes("00"); // 00 in flight
        appendBytes("08a"); // 0008a in flight
        appendBytes("bcd00"); // 00 in flight
        appendBytes("00"); // no packets in flight

        appendBytes("000"); // 000 in flight
        appendBytes("5aPA"); // PA in flight
        appendBytesAndClose("CKhellothere");

        assertPacketLines("0009abcde", "0000", "0006aa");
        assertPackEquals("REPLACEDhellothere");
        assertFalse(filter.invalid);
        assertEquals("0009abcde00000006aaREPLACEDhellothere", delegate.getData());
        assertTrue(delegate.closed);
    }

    @Test
    public void testInvalidGitPacketStream() {
        appendBytes("00000"); //0 in flight
        appendBytes("008abc"); // 0008abc in flight
        appendBytes("dthis is not a packet line");
        appendBytesAndClose("this is not a packet line either");

        //Nothing should have been sent to the filter past the point where the stream became invalid
        assertPacketLines("0000", "0008abcd");
        assertPackEquals("");
        assertTrue(filter.invalid);

        //The delegate should still receive all data
        assertEquals("00000008abcdthis is not a packet linethis is not a packet line either", delegate.getData());
        assertTrue(delegate.closed);

        logListener.assertWarn("Unexpected packet line prefix");
    }

    @Test
    public void testFiltersAndDelegateCanOptOut() {
        //Mock a filter that opts out after the second packet line and first pack chunk
        MutableInt packetLineCount = new MutableInt();
        MutableInt packCount = new MutableInt();
        GitReceivePackFilter optOut = mock(GitReceivePackFilter.class);
        when(optOut.wantsPacketLines()).thenAnswer(ignored -> packetLineCount.intValue() < 2);
        when(optOut.wantsPack()).thenAnswer(ignored -> packCount.intValue() == 0);
        when(optOut.onPack(any())).thenAnswer(invocation -> {
            packCount.increment();
            return invocation.getArgument(0);
        });
        when(optOut.onPacketLine(any())).thenAnswer(invocation -> {
            packetLineCount.increment();
            return invocation.getArgument(0);
        });

        //Mock a delegate that opts out after 6 chunks (4 packet lines + 2 PACK chunks)
        MutableInt count = new MutableInt();
        SatiableConsumer<RpcStreamChunk> delegate = mock(SatiableConsumer.class);
        when(delegate.accept(any())).thenAnswer(ignored -> count.incrementAndGet() >= 6);

        stream = new FilteringGitReceivePackStream(delegate);
        stream.addFilter(optOut);
        stream.addFilter(filter);

        appendBytes("0005a0006ab"); // write two packet lines
        appendBytes("0007abc0"); // write another packet line + start of next packet line
        appendBytes("000PACK\0this is the pack\n"); // packet line + first PACK chunk
        //3 more PACK chunks
        appendBytes("more pack bytes\n");
        appendBytes("even more pack bytes\n");
        appendBytesAndClose("ever more pack bytes\n");

        //Verify that the optOut filter received exactly the expected amount of callbacks
        assertEquals(2, packetLineCount.intValue());
        assertEquals(1, packCount.intValue());

        //Verify that the delegate only received 6 chunks
        verify(delegate, times(6)).accept(any());

        //Verify that the other filter (CapturingFilter) received all chunks sent to the delegate
        assertPacketLines("0005a", "0006ab", "0007abc", "0000");
        assertPackEquals("PACK\0this is the pack\nmore pack bytes\n");
    }

    @Test
    public void testPartialPacketIsOutputOnClose() {
        appendBytesAndClose("0008abcd00"); // 00 in flight

        assertPacketLines("0008abcd");
        assertEquals("0008abcd00", delegate.getData());
        assertTrue(delegate.closed);
    }

    @Test
    public void testPartialPackets() {
        appendBytes("00"); // 00 in flight
        appendBytes("08a"); // 0008a in flight
        appendBytes("bcd00"); // 00 in flight
        appendBytes("0"); // 000 in flight - this chunk does not complete a packet header
        appendBytes("0"); // no packets in flight

        appendBytes("000"); // 000 in flight
        appendBytes("5a0009ab"); // 0009ab in flight
        appendBytesAndClose("cdePACKhellothere");

        assertPacketLines("0008abcd", "0000", "0005a", "0009abcde");
        assertPackEquals("PACKhellothere");
        assertFalse(filter.invalid);
        assertEquals("0008abcd00000005a0009abcdePACKhellothere", delegate.getData());
        assertTrue(delegate.closed);
    }

    @Test
    public void testPartialPacketsSplitAcrossBuffers() {
        appendBytes("00", "08a"); // 00, then 008a in flight
        appendBytes("bcd00", "00", "0"); // 00, then nothing, then 0 in flight

        appendBytesAndClose("00", "5aPAC", "Khello"); // 000, then PAC in flight, then in pack section

        assertPackEquals("PACKhello");
        assertTrue(delegate.closed);
    }

    @Test
    public void testReceivePackRealResponse() throws IOException {
        byte[] buffer = new byte[17];
        try (InputStream in = new ClassPathResource("git-receive-pack.response", getClass()).getInputStream()) {
            int read;
            while ((read = in.read(buffer)) != -1) {
                appendBytes(buffer, 0, read);
            }
            appendBytesAndClose();
        }

        assertPacketLines("00968621398e3d405821832d0830fc53c560d427b9fb d98685844c5146e89b2eb1c734ad788a0d8501ca " +
                        "refs/heads/master\0 report-status side-band-64k agent=git/2.18.0\n",
                "0000");
        assertPackEquals("PACK0000");
        assertFalse(filter.invalid);
        assertEquals("00968621398e3d405821832d0830fc53c560d427b9fb d98685844c5146e89b2eb1c734ad788a0d8501ca " +
                        "refs/heads/master\0 report-status side-band-64k agent=git/2.18.0\n0000PACK0000",
                delegate.getData());
        assertTrue(delegate.closed);
    }

    @Test
    public void testReceivePackRealStdin() throws IOException {
        byte[] buffer = new byte[16384];
        try (InputStream in = new ClassPathResource("git-receive-pack-medium.stdin", getClass()).getInputStream()) {
            int read;
            while ((read = in.read(buffer)) != -1) {
                appendBytes(buffer, 0, read);
            }
            appendBytesAndClose();
        }

        List<String> packetLines = filter.getPacketLines();
        assertEquals(312, packetLines.size());
        assertEquals("006b0000000000000000000000000000000000000000 f82c1cb9af618e1b689872f53745b562efc536d2 " +
                "refs/heads/hahaahaha2", packetLines.get(1));
        assertEquals("007a0000000000000000000000000000000000000000 a884f74010f4ce493daaede11c7ce5ed35a558a2 " +
                "refs/remotes/origin/test1/branch-500", packetLines.get(310));
        assertEquals("0000", packetLines.get(311));
        assertPackEquals("PACK\0\0\0\2\0\0\0=");
        assertFalse(filter.invalid);
        assertTrue(delegate.closed);
    }

    @Test
    public void testWithoutFilters() {
        stream = new FilteringGitReceivePackStream(delegate);

        appendBytes("0006aa000000"); // 0006aa, 0000 and partial 0 packet
        assertEquals("0006aa000000", delegate.getData());

        appendBytesAndClose("008abcd0000PACKtesting");
        assertEquals("0006aa000000008abcd0000PACKtesting", delegate.getData());
    }

    @Test
    public void testWithoutPacketLineFilters() {
        stream = new FilteringGitReceivePackStream(delegate);
        GitReceivePackFilter filter = mock(GitReceivePackFilter.class);
        when(filter.wantsPack()).thenReturn(true);
        when(filter.onPack(any())).thenAnswer(returnFirst());
        stream.addFilter(filter);

        appendBytes("0006aa000000"); // 0006aa, 0000 and partial 00 packet
        assertEquals("0006aa0000", delegate.getData());

        appendBytes("08abcd00"); // finished the 008abcd packet, partial 00 packet
        assertEquals("0006aa00000008abcd", delegate.getData());

        appendBytes("10abc"); // partial 0010abc packet in flight
        //A partial packet line is in flight, but since the packet line header was read, it can already be output
        assertEquals("0006aa00000008abcd0010abc", delegate.getData());

        appendBytes("defghijklPA"); // finished 010abcdefghijkl packet, started PACK header
        assertEquals("0006aa00000008abcd0010abcdefghijkl", delegate.getData());

        appendBytes("CKtesting\n");
        assertEquals("0006aa00000008abcd0010abcdefghijklPACKtesting\n", delegate.getData());
        when(filter.wantsPack()).thenReturn(false);

        appendBytes("more pack data\n");
        appendBytesAndClose("more pack data\n");

        assertEquals("0006aa00000008abcd0010abcdefghijklPACKtesting\nmore pack data\nmore pack data\n",
                delegate.getData());
    }

    private void appendBytes(byte[] buffer, int offset, int length) {
        stream.write(RpcStreamChunk.newBuilder()
                .setData(ByteString.copyFrom(buffer, offset, length))
                .build());
    }

    private void appendBytes(String... fragments) {
        appendChunk(false, fragments);
    }

    private void appendBytesAndClose(String... fragments) {
        appendChunk(true, fragments);
    }

    private void appendChunk(boolean closed, String... fragments) {
        ByteString data = null;
        for (String fragment : fragments) {
            ByteString fragmentString = ByteString.copyFromUtf8(fragment);
            data = data == null ? fragmentString : data.concat(fragmentString);
        }

        stream.write(RpcStreamChunk.newBuilder()
                .setClosed(closed)
                .setData(data == null ? ByteString.EMPTY : data)
                .build());
    }

    private void assertPackEquals(String expected) {
        assertEquals(expected, new String(filter.pack.toByteArray(), UTF_8));
    }

    private void assertPacketLines(String... expected) {
        assertEquals(ImmutableList.copyOf(expected), filter.getPacketLines());
    }

    private static class CapturingFilter implements GitReceivePackFilter {

        private final ByteArrayOutputStream pack = new ByteArrayOutputStream();
        private final List<RpcStreamChunk> packetLines = new ArrayList<>();

        private boolean invalid;

        @Override
        public void onInvalid() {
            invalid = true;
        }

        @Nonnull
        @Override
        public RpcStreamChunk onPack(@Nonnull RpcStreamChunk chunk) {
            try {
                chunk.getData().writeTo(pack);
            } catch (IOException e) {
                throw new UncheckedIOException(e);
            }
            return chunk;
        }

        @Nonnull
        @Override
        public RpcStreamChunk onPacketLine(@Nonnull RpcStreamChunk chunk) {
            packetLines.add(chunk);
            return chunk;
        }

        @Override
        public boolean wantsPack() {
            return true;
        }

        @Override
        public boolean wantsPacketLines() {
            return true;
        }

        List<String> getPacketLines() {
            return packetLines.stream()
                    .map(chunk -> chunk.getData().toStringUtf8())
                    .collect(Collectors.toList());
        }
    }

    private static class CapturingRpcStream implements SatiableConsumer<RpcStreamChunk> {

        private final ByteArrayOutputStream data = new ByteArrayOutputStream();

        private boolean closed;

        @Override
        public boolean accept(RpcStreamChunk value) {
            try {
                assertFalse("Received a chunk after a previous chunk closed the stream", closed);
                value.getData().writeTo(data);
                closed = value.getClosed();
            } catch (IOException e) {
                throw new UncheckedIOException(e);
            }
            return false;
        }

        public String getData() {
            return new String(data.toByteArray(), UTF_8);
        }
    }

    private static class MutatingFilter implements GitReceivePackFilter {

        private final ByteString packReplacement;
        private final Map<ByteString, ByteString> packetLineTransforms;
        private boolean packModified;

        MutatingFilter(String packReplacement, Map<String, String> packetLineTransforms) {
            this.packetLineTransforms = packetLineTransforms.entrySet().stream()
                    .collect(Collectors.toMap(
                            entry -> ByteString.copyFromUtf8(entry.getKey()),
                            entry -> ByteString.copyFromUtf8(entry.getValue())));
            this.packReplacement = ByteString.copyFromUtf8(packReplacement);
        }

        @Nonnull
        @Override
        public RpcStreamChunk onPack(@Nonnull RpcStreamChunk chunk) {
            if (packModified) {
                return chunk;
            }

            packModified = true;
            return chunk.toBuilder()
                    .setData(packReplacement.concat(chunk.getData().substring(4)))
                    .build();
        }

        @Nonnull
        @Override
        public RpcStreamChunk onPacketLine(@Nonnull RpcStreamChunk chunk) {
            ByteString replacement = packetLineTransforms.get(chunk.getData());
            return replacement == null ? chunk : chunk.toBuilder()
                    .setData(replacement)
                    .build();
        }

        @Override
        public boolean wantsPack() {
            return true;
        }

        @Override
        public boolean wantsPacketLines() {
            return true;
        }
    }
}
