package com.atlassian.bitbucket.mesh.process.nu;

import com.atlassian.bitbucket.mesh.execution.ExecutionManager;
import com.atlassian.bitbucket.mesh.process.*;
import com.atlassian.bitbucket.mesh.process.handler.AbstractLineStdioHandler;
import com.atlassian.bitbucket.mesh.process.handler.AbstractProcessStdioHandler;
import com.atlassian.bitbucket.mesh.throttle.ThrottleService;
import com.atlassian.bitbucket.mesh.throttle.ThrottledResource;
import com.atlassian.bitbucket.mesh.throttle.Ticket;
import com.atlassian.bitbucket.mesh.util.Version;
import de.schlichtherle.truezip.file.TArchiveDetector;
import de.schlichtherle.truezip.file.TFile;
import org.apache.commons.lang3.SystemUtils;
import org.junit.*;
import org.junit.rules.ExpectedException;
import org.junit.rules.TemporaryFolder;
import org.junit.runner.RunWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;
import org.springframework.core.io.ClassPathResource;

import javax.annotation.Nonnull;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.nio.file.Path;
import java.time.Duration;
import java.util.concurrent.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import static com.atlassian.bitbucket.mesh.throttle.ThrottledResource.COMMAND;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.*;
import static org.junit.Assert.*;
import static org.junit.Assume.assumeTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.same;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class NuProcessCreatorTest {

    @ClassRule
    public static final TemporaryFolder tempDir = new TemporaryFolder();

    private static Path repositoryDir;

    @Rule
    public final ExpectedException thrown = ExpectedException.none();

    @Mock
    private ProcessCompletionHandler completionHandler;
    @Mock
    private ProcessConfigurer configurer;
    @Mock
    private ExecutionManager executionManager;
    private NuProcessCreator processCreator;
    @Mock
    private ThrottleService throttleService;
    @Mock
    private Ticket ticket;

    //In order to make the Git operations performed by this test more reliable, we need a fixed data
    //set for them to operate on. This is also important for the source build which doesn't run in a
    //clone of the codebase and therefore can't access its history.
    //
    //Since none of the Git operations this test uses write to the repository, we just unpack it once,
    //before the first test, and we use a TemporaryDirectory @ClassRule to cleanup after the last one.
    @BeforeClass
    public static void prepareRepository() throws Exception {
        repositoryDir = tempDir.newFolder("repository").toPath();

        ClassPathResource repository = new ClassPathResource("git/default-repository.zip");
        TFile.cp_r(new TFile(repository.getFile()), repositoryDir.toFile(), TArchiveDetector.ALL);
    }

    @Before
    public void setup() throws Exception {
        when(executionManager.maybeWrap(any(Runnable.class))).thenAnswer(invocation -> invocation.getArgument(0));
        when(throttleService.acquireTicket(any())).thenReturn(ticket);

        processCreator = new NuProcessCreator(new ProcessConfigurer[]{configurer}, executionManager, throttleService);

        doAnswer(invocation -> {
            ProcessSummary summary = invocation.getArgument(0);
            if (summary.getThrown() != null) {
                throw summary.getThrown();
            }

            return null;
        }).when(completionHandler).onComplete(any());
    }

    @After
    public void teardown() {
        processCreator.shutdown();
    }

    @Ignore("For some reason this test is really flaky on CI")
    @Test(timeout = 5000L)
    public void testCancelFuture() throws Exception {
        CountDownLatch exited = new CountDownLatch(1);
        CountDownLatch running = new CountDownLatch(1);

        ProcessParameters<Void> parameters = heartbeat(new AbstractProcessStdioHandler<Void>() {

            private boolean started;

            @Override
            public void onExit(int exitCode) {
                if (!started) {
                    //If the process fails to start (nonexistent binary?), ensure the test doesn't
                    //hang indefinitely waiting for the running latch to be released
                    running.countDown();
                }
                exited.countDown();
            }

            @Override
            public Void getResult() {
                return null;
            }

            @Override
            public void onStart(@Nonnull OsProcess process) {
                super.onStart(process);

                started = true;
            }

            @Override
            public void onStdout(@Nonnull ByteBuffer buffer, boolean closed) {
                buffer.position(buffer.position() + buffer.remaining());

                //Upon receiving the first line of stdout, allow the test to cancel the process
                running.countDown();
            }
        });
        Future<Void> future = processCreator.start(parameters);

        //Wait for the process to start
        running.await();

        try {
            //Verify the process is still running
            future.get(0, TimeUnit.MILLISECONDS);

            fail("The process finished on its own before it could be canceled");
        } catch (TimeoutException e) {
            //The process didn't complete on its own. Now we cancel it
            assertTrue(future.cancel(true));
            assertTrue(future.isCancelled());

            assertTrue("The process did not terminate after canceling", exited.await(5L, TimeUnit.SECONDS));
        }
    }

    @Test
    public void testCancelProcess() {
        ProcessParameters<String> parameters = git(new AbstractLineStdioHandler<String>(StandardCharsets.UTF_8) {

            private String line;

            @Override
            public String getResult() {
                return line;
            }

            @Override
            protected boolean onStdout(@Nonnull String line, boolean truncated) {
                this.line = line;

                return false; //Returning false will cancel the process after the first line
            }
        }, "log", "--decorate", "-1", "HEAD");

        String line = processCreator.run(parameters);
        assertNotNull(line);
        assertThat(line, containsString("HEAD ->"));
    }

    @Test
    public void testDestroyWaitingProcess() {
        RuntimeException exception = new RuntimeException("Boom");

        thrown.expect(sameInstance(exception));

        //git cat-file --batch blocks waiting for stdin and doesn't exit until it's closed. This test
        //verifies such a process is correctly destroyed after a handler failure
        CountDownLatch terminated = new CountDownLatch(1);
        ProcessParameters<Void> parameters = git(new AbstractProcessStdioHandler<Void>() {

            @Override
            public Void getResult() {
                return null;
            }

            @Override
            public void onExit(int exitCode) {
                terminated.countDown();
            }

            @Override
            public void onStart(@Nonnull OsProcess process) {
                super.onStart(process);

                process.wantWrite();
            }

            @Override
            public boolean onStdinReady(@Nonnull ByteBuffer buffer) {
                buffer.put("7e164127cf77b202a6c4269b1bd43b91fd9543bf\n".getBytes(StandardCharsets.UTF_8));
                buffer.flip();

                return false;
            }

            @Override
            public void onStdout(@Nonnull ByteBuffer buffer, boolean closed) {
                //This will trigger the process to be destroyed
                throw exception;
            }
        }, "cat-file", "--batch");

        Future<Void> future = processCreator.start(parameters);
        try {
            future.get(5L, TimeUnit.SECONDS);
        } catch (ExecutionException e) {
            throw (RuntimeException) e.getCause();
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        } catch (TimeoutException e) {
            //This means the handling is broken. We try cancel as a last-ditch effort just to prevent
            //any zombie process. Note that even if this does bring the process down, the test will
            //still (correctly) fail because the expected exception isn't thrown
            System.out.println("git cat-file --batch didn't terminate; trying to cancel");
            future.cancel(true);
        } finally {
            try {
                assumeTrue("git cat-file --batch never terminated", terminated.await(1L, TimeUnit.SECONDS));
            } catch (InterruptedException ex) {
                Thread.currentThread().interrupt();
            }
        }
    }

    @Test
    public void testExecutionContextIsNotConfiguredWhenRun() {
        //When a Process is run the I/O is pumped on the current request thread, so the current ExecutionContext
        //is already set up correctly.
        processCreator.run(gitVersion(new AbstractProcessStdioHandler<Void>() {

            @Override
            public Void getResult() {
                return null;
            }

            @Override
            public void onStdout(@Nonnull ByteBuffer buffer, boolean closed) {
                buffer.position(buffer.position() + buffer.remaining());
            }
        }));

        verifyNoInteractions(executionManager);
    }

    @Test
    public void testExecutionContextIsConfiguredWhenStarted() throws Exception {
        Future<Void> future = processCreator.start(gitVersion(new AbstractProcessStdioHandler<Void>() {

            @Override
            public Void getResult() {
                return null;
            }

            @Override
            public void onStdout(@Nonnull ByteBuffer buffer, boolean closed) {
                buffer.position(buffer.position() + buffer.remaining());
            }
        }));
        try {
            future.get(15L, TimeUnit.SECONDS);
        } catch (ExecutionException e) {
            throw (Exception) e.getCause();
        } finally {
            if (!future.isDone()) {
                future.cancel(true);
            }
        }
        verify(executionManager).maybeWrap(any(Runnable.class));
    }

    @Ignore("This test is too slow to run by default")
    @Test
    public void testIdleTimeout() throws Exception {
        ProcessStdioHandler<Void> handler = new AbstractProcessStdioHandler<Void>() {

            @Override
            public Void getResult() {
                return null;
            }

            @Override
            public void onStdout(@Nonnull ByteBuffer buffer, boolean closed) {
                buffer.position(buffer.position() + buffer.remaining());
            }
        };
        ProcessParameters<Void> processParameters = new ProcessParameters.Builder<>(handler, completionHandler)
                .arguments("cat")
                .executionInterval(Duration.ofSeconds(5L))
                .idleInterval(Duration.ofSeconds(3L))
                .build();

        Future<Void> future = processCreator.start(processParameters);
        try {
            future.get(15L, TimeUnit.SECONDS);
        } catch (ExecutionException e) {
            throw (Exception) e.getCause();
        } finally {
            if (!future.isDone()) {
                future.cancel(true);
            }
        }

        ArgumentCaptor<ProcessSummary> summaryCaptor = ArgumentCaptor.forClass(ProcessSummary.class);
        verify(completionHandler).onComplete(summaryCaptor.capture());

        ProcessSummary summary = summaryCaptor.getValue();
        assertTrue("The process timeout should have been enforced", summary.isTimedOut());
        assertThat("Timeout should trigger a non-zero exit", summary.getExitCode(), not(equalTo(0)));
        assertNull("No blanket exception should be created", summary.getThrown());
    }

    @Test
    public void testOnStdoutExceptionHandling() {
        RuntimeException exception = new RuntimeException("Boom");

        thrown.expect(sameInstance(exception));

        processCreator.run(gitVersion(new AbstractProcessStdioHandler<Void>() {

            @Override
            public Void getResult() {
                return null;
            }

            @Override
            public void onStdout(@Nonnull ByteBuffer buffer, boolean closed) {
                throw exception;
            }
        }));
    }

    @Test
    public void testOutputParsingForRun() {
        Integer lines = processCreator.run(git(new LineCountStdioHandler(), "diff", "-C", "HEAD^", "HEAD"));
        assertNotNull(lines);
        assertThat(lines, greaterThan(0));
    }

    @Test
    public void testOutputParsingForStart() throws Exception {
        Future<Integer> future = processCreator.start(git(new LineCountStdioHandler(), "diff", "-C", "HEAD^", "HEAD"));

        Integer lines = future.get();
        assertNotNull(lines);
        assertThat(lines, greaterThan(0));
    }

    @Test
    public void testRun() {
        ProcessParameters<Version> parameters = gitVersion(new VersionStdioHandler());

        Version version = processCreator.run(parameters);
        assertNotNull(version);

        verify(configurer).configure(same(parameters));
        verifyNoInteractions(throttleService);
    }

    @Test
    public void testRunAcquiresTicket() {
        ProcessParameters<Version> parameters = gitVersion(new VersionStdioHandler(), COMMAND);

        Version version = processCreator.run(parameters);
        assertNotNull(version);

        verify(configurer).configure(same(parameters));
        verify(throttleService).acquireTicket(same(COMMAND));
        verify(ticket).close();
    }

    @Test
    public void testRunThrowsOnTicketExhaustion() {
        //At runtime acquireTicket will actually throw a ResourceBusyException, but with the current Maven
        //dependencies it's not possible to do that for this test. The specific exception type here isn't
        //actually important for this test, though; what this test is really validating is that whatever
        //exception acquireTicket throws gets propagated, which can be tested with any exception
        //TODO: Fix this exception when the repositories are split
        thrown.expect(RejectedExecutionException.class);

        doThrow(RejectedExecutionException.class).when(throttleService).acquireTicket(any());

        ProcessStdioHandler<?> handler = mock(ProcessStdioHandler.class);
        try {
            processCreator.run(gitVersion(handler, COMMAND));
        } finally {
            verify(throttleService).acquireTicket(same(COMMAND));
            verifyNoInteractions(handler);
        }
    }

    @Test
    public void testStart() throws Exception {
        ProcessParameters<Version> parameters = gitVersion(new VersionStdioHandler());

        Future<Version> future = processCreator.start(parameters);
        Version version = future.get();
        assertNotNull(version);

        verify(configurer).configure(same(parameters));
        verifyNoInteractions(throttleService);
    }

    @Test
    public void testStartAcquiresTicket() throws Exception {
        ProcessParameters<Version> parameters = gitVersion(new VersionStdioHandler(), COMMAND);

        Future<Version> future = processCreator.start(parameters);
        Version version = future.get();
        assertNotNull(version);

        verify(configurer).configure(same(parameters));
        verify(throttleService).acquireTicket(same(COMMAND));
        //the ticket is closed after the future completes, so use a timeout
        verify(ticket, timeout(150L)).close();
    }

    @Test
    public void testStartThrowsOnTicketExhaustion() throws Exception {
        //See testRunThrowsOnTicketExhaustion for an explanation of this exception
        thrown.expect(RejectedExecutionException.class);

        doThrow(RejectedExecutionException.class).when(throttleService).acquireTicket(any());

        ProcessStdioHandler<?> handler = mock(ProcessStdioHandler.class);
        try {
            processCreator.start(gitVersion(handler, COMMAND)).get();
        } finally {
            verify(throttleService).acquireTicket(same(COMMAND));
            verifyNoInteractions(handler);
        }
    }

    private <T> ProcessParameters<T> git(ProcessStdioHandler<T> stdioHandler, String... arguments) {
        return new ProcessParameters.Builder<>(stdioHandler, completionHandler)
                .arguments("git", arguments)
                .workDir(repositoryDir)
                .build();
    }

    private <T> ProcessParameters<T> gitVersion(ProcessStdioHandler<T> stdioHandler) {
        return gitVersion(stdioHandler, null);
    }

    private <T> ProcessParameters<T> gitVersion(ProcessStdioHandler<T> stdioHandler, ThrottledResource resource) {
        return new ProcessParameters.Builder<>(stdioHandler, completionHandler)
                .arguments("git", "version")
                .throttled(resource)
                .build();
    }

    /**
     * Builds a command that, when executed, will run until it's killed and produce steady output. The goal here
     * is not to produce large amounts of output; rather, it's to produce a small amount of output with a regular
     * cadence.
     *
     * @return parameters for a command that will produce output until it's killed
     */
    private <T> ProcessParameters<T> heartbeat(ProcessStdioHandler<T> stdioHandler) {
        ProcessParameters.Builder<T> builder = new ProcessParameters.Builder<>(stdioHandler, completionHandler);
        if (SystemUtils.IS_OS_MAC) {
            builder.arguments("ping", "localhost");
        } else {
            //Note: The Linux PBC image used by Bamboo doesn't have ping, so this uses watch instead
            builder.arguments("watch", "-n", "1", "-x", "date");
        }

        return builder.build();
    }

    private static class LineCountStdioHandler extends AbstractLineStdioHandler<Integer> {

        private int lines;

        LineCountStdioHandler() {
            super(StandardCharsets.UTF_8);
        }

        @Override
        public Integer getResult() {
            return lines;
        }

        @Override
        protected boolean onStdout(@Nonnull String line, boolean truncated) {
            ++lines;

            return true;
        }
    }

    private static class VersionStdioHandler extends AbstractLineStdioHandler<Version> {

        private Version version;

        VersionStdioHandler() {
            super(StandardCharsets.UTF_8);
        }

        @Override
        public Version getResult() {
            return version;
        }

        @Override
        public boolean onStdout(@Nonnull String line, boolean truncated) {
            if (version == null) {
                Matcher matcher = Pattern.compile("^git version (.*)").matcher(line);
                if (matcher.matches()) {
                    version = new Version(matcher.group(1));

                    return false;
                }
            }

            return true;
        }
    }
}
