package com.atlassian.bitbucket.scm.cache.ssh;

import com.atlassian.bitbucket.scm.cache.TestSupport;
import com.atlassian.bitbucket.scm.git.protocol.ssh.GitSshScmRequestHandler;
import com.atlassian.bitbucket.scm.ssh.ExitCodeCallback;
import com.atlassian.bitbucket.scm.ssh.SshScmRequest;
import com.atlassian.bitbucket.throttle.ThrottleService;
import com.atlassian.bitbucket.throttle.Ticket;
import com.atlassian.bitbucket.util.CancelState;
import com.google.common.base.Charsets;
import com.google.common.io.ByteStreams;
import org.apache.commons.io.IOUtils;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.runners.MockitoJUnitRunner;

import javax.annotation.Nonnull;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Optional;
import java.util.concurrent.*;

import static java.util.Optional.of;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.any;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class UploadPackProxyTest {

    public static final String SCM_HOSTING = "scm-hosting";
    @Mock
    private CancelState cancelState;
    @Mock
    private ThrottleService throttleService;
    @Mock
    private Ticket ticket;
    private ExecutorService executor;
    private TestSupport.MockProcess process;
    private Exception failure;

    private GitSshScmRequestHandler requestHandler;
    private UploadPackProxy proxy;

    @After
    public void teardown() throws Exception {
        executor.shutdownNow();
        assertTrue("Failed to cleanup thread-pool", executor.awaitTermination(5, TimeUnit.SECONDS));
        executor = null;
        if (failure != null) {
            throw failure;
        }
    }

    @Before
    public void setup() {
        executor = Executors.newCachedThreadPool();
        failure = null;
        process = new TestSupport.MockProcess();
        requestHandler = new GitSshScmRequestHandler() {
            @Nonnull
            @Override
            public Optional<SshScmRequest> create(@Nonnull String s, @Nonnull InputStream in, @Nonnull OutputStream out,
                                        @Nonnull OutputStream err, @Nonnull ExitCodeCallback exitCodeCallback) {

                SshScmRequest request = mock(SshScmRequest.class);
                try {
                    doAnswer(invocationOnMock -> {
                        // Client input is pumped directly to the SCM process, but we don't necessarily need the
                        // client to close its end of the connection, so we don't signal EOF or wait for it.
                        Future<?> future = pump(in, process.getStdIn());

                        // Output and errors from the SCM process are pumped directly, and EOF signaled so that we
                        // can wait for both streams to be fully flushed.
                        CountDownLatch latch = new CountDownLatch(2);
                        pumpAndSignalEOF(process.getStdOut(), out, latch);
                        pumpAndSignalEOF(process.getStdErr(), err, latch);

                        latch.await();
                        future.cancel(true);
                        return null;
                    }).when(request).handleRequest();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
                return of(request);
            }

            @Override
            public boolean supports(@Nonnull String remoteCommand) {
                return true;
            }
        };
        when(throttleService.acquireTicket(SCM_HOSTING)).thenReturn(mock(Ticket.class));
        when(throttleService.acquireTicket(SCM_HOSTING)).thenReturn(ticket);

        proxy = new UploadPackProxy(executor, "bla", cancelState, requestHandler, throttleService);
    }

    @Test
    public void testRefsOnly() throws Exception {
        process.output(stream("ref.advertisement"));
        ByteArrayOutputStream output = new ByteArrayOutputStream();

        // read the refs from the process
        proxy.streamRefs(output);
        proxy.awaitDelegateScmRequestCompletion(5, TimeUnit.SECONDS);

        assertEquals(IOUtils.toString(stream("ref.advertisement")), output.toString(Charsets.UTF_8.name()));
    }

    @Test
    public void ticketAcquired() throws Exception {
        process.output(stream("ref.advertisement"));

        proxy.streamRefs(ByteStreams.nullOutputStream());
        proxy.awaitDelegateScmRequestCompletion(5, TimeUnit.SECONDS);

        verify(throttleService).acquireTicket(SCM_HOSTING);
    }

    @Test
    public void ticketClosedOnceNormalOperation() throws Exception {
        process.output(stream("ref.advertisement"));

        proxy.streamRefs(ByteStreams.nullOutputStream());
        proxy.awaitDelegateScmRequestCompletion(5, TimeUnit.SECONDS);
        proxy.releaseTicket();

        verify(ticket).close();
    }

    @Test(expected = IllegalStateException.class)
    public void ticketClosedIfExecutorRejectsTask() throws Exception {
        ExecutorService mockExecutor = mock(ExecutorService.class);
        when(mockExecutor.submit(any(FutureTask.class))).thenThrow(new RejectedExecutionException());
        UploadPackProxy proxy = new UploadPackProxy(mockExecutor, "bla", cancelState, requestHandler, throttleService);

        process.output(stream("ref.advertisement"));

        try {
            proxy.streamRefs(ByteStreams.nullOutputStream());
        } finally {
            verify(ticket).close();
        }
    }

    private Future<?> pump(InputStream inputStream, OutputStream outputStream) {
        return executor.submit(() -> {
            try {
                IOUtils.copy(inputStream, outputStream);
            } catch (Throwable t) {
                // There is a race condition under which the task spawned by
                // UploadPackProxy.startDelegateRequest() can close the
                // InputStream before this thread has started causing an IOException to be thrown.
                // We ignore IOException here because none of the current tests stream data to the process input.
                if (!(t instanceof IOException)) {
                    failInJunitThread("failed during IO pump ", t);
                }
            }
        });
    }

    private void pumpAndSignalEOF(InputStream inputStream, OutputStream outputStream, CountDownLatch latch) {
        executor.submit(() -> {
            try {
                IOUtils.copy(inputStream, outputStream);
            } catch (Throwable t) {
                failInJunitThread("failed during IO pump", t);
            } finally {
                latch.countDown();
            }
        });
    }

    private void failInJunitThread(String message, Throwable t) {
        failure = new Exception(message, t);
    }

    private InputStream stream(String fileName) {
        return getClass().getResourceAsStream("/git/" + fileName);
    }
}
