package acceptance;

import com.google.common.io.Resources;
import io.digdag.client.DigdagClient;
import io.digdag.client.api.Id;
import io.digdag.client.api.RestSessionAttempt;
import io.digdag.client.api.RestTask;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import utils.CommandStatus;
import utils.TemporaryDigdagServer;
import utils.TestUtils;

import java.net.ConnectException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.time.Duration;
import javax.ws.rs.ProcessingException;
import javax.ws.rs.ServiceUnavailableException;

import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.Arrays.asList;
import static utils.TestUtils.getAttemptId;
import static utils.TestUtils.main;
import static org.hamcrest.Matchers.containsString;
import static org.hamcrest.Matchers.greaterThan;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.instanceOf;
import static org.junit.Assert.assertThat;
import static org.junit.Assume.assumeThat;

public class ServerGracefulShutdownIT
{
    private final Logger logger = LoggerFactory.getLogger(ServerGracefulShutdownIT.class);

    @Rule
    public TemporaryFolder folder = new TemporaryFolder();

    @Rule
    public TemporaryDigdagServer server = TemporaryDigdagServer.builder()
            .inProcess(false)
            .build();

    private Path root()
    {
        return folder.getRoot().toPath().toAbsolutePath();
    }

    private Path config;
    private Path projectDir;
    private DigdagClient client;

    @Before
    public void setUp()
            throws Exception
    {
        assumeThat(server.hasUnixProcess(), is(true));

        projectDir = root().resolve("foobar");
        config = folder.newFile().toPath();

        client = DigdagClient.builder()
                .host(server.host())
                .port(server.port())
                .build();
    }

    private Id startSleepTask()
        throws Exception
    {
        Id attemptId;
        {
            CommandStatus initStatus = main("init",
                    "-c", config.toString(),
                    projectDir.toString());
            assertThat(initStatus.code(), is(0));

            Files.write(projectDir.resolve("sleep.dig"), asList(Resources.toString(
                    Resources.getResource("acceptance/server_graceful_shutdown/sleep.dig"), UTF_8)
                    .replace("${outdir}", root().toString())));

            CommandStatus pushStatus = main("push",
                    "--project", projectDir.toString(),
                    "server_graceful_shutdown",
                    "-c", config.toString(),
                    "-e", server.endpoint());
            assertThat(pushStatus.errUtf8(), pushStatus.code(), is(0));

            CommandStatus startStatus = main("start",
                    "-c", config.toString(),
                    "-e", server.endpoint(),
                    "server_graceful_shutdown", "sleep",
                    "--session", "now");
            assertThat(startStatus.code(), is(0));
            attemptId = getAttemptId(startStatus);
        }

        // Wait for the task to start
        TestUtils.expect(Duration.ofMinutes(5), () -> {
            RestTask checkerTask = client.getTasks(attemptId).getTasks()
                .stream()
                .filter(it -> it.getFullName().endsWith("+start_checker"))
                .findFirst()
                .orElseThrow(() -> new IllegalStateException("+start_checker task doesn't exist"));

            if (checkerTask.getState().equals("success")) {
                return true;
            }

            RestSessionAttempt attempt = client.getSessionAttempt(attemptId);
            if (attempt.getDone()) {
                logger.warn("The attempt itself already finished... Is it expected situation....? attempt={}", attempt);
                return true;
            }

            return false;
        });

        return attemptId;
    }

    @Test
    public void gracefulShutdown()
            throws Exception
    {
        Id attemptId = startSleepTask();

        server.terminateProcess();

        // server started termination but it should be alive at most 5 seconds.

        int aliveSeconds = 0;
        while (true) {
            try {
                client.getSessionAttempt(attemptId);
                aliveSeconds++;
            }
            catch (Exception ex) {
                // if REST API fails, the cause should be 503 Service Unavailable or
                // connection refused.
                if (ex instanceof ProcessingException) {
                    assertThat(ex.getCause(), instanceOf(ConnectException.class));
                    break;
                }
                else {
                    assertThat(ex, instanceOf(ServiceUnavailableException.class));
                    break;
                }
            }

            if (aliveSeconds > Duration.ofMinutes(5).getSeconds()) {
                throw new IllegalStateException("Server doesn't shutdown");
            }

            Thread.sleep(1000);
        }

        // all running tasks should be done
        assertThat(Files.exists(root().resolve("done.out")), is(true));

        // but waiting tasks should not start
        assertThat(Files.exists(root().resolve("after_sleep.out")), is(false));

        // REST API should be alive for a while
        assertThat(aliveSeconds, greaterThan(5));

        assertThat(server.outUtf8(), containsString("Waiting for completion of 2 running tasks..."));
        assertThat(server.outUtf8(), containsString("Closing HTTP listening sockets"));

        TestUtils.expect(Duration.ofMinutes(5), () -> !server.isProcessAlive());

        assertThat(server.outUtf8(), containsString("Shutting down HTTP worker threads"));
        assertThat(server.outUtf8(), containsString("Shutting down system"));
        assertThat(server.outUtf8(), containsString("Shutdown completed"));
    }
}

@@ -578,24 +578,50 @@
 e.Duration;%0A
+import java.time.Instant;%0A
 import javax
@@ -4256,16 +4256,68 @@
 ocess();
+%0A        Instant terminateStartedAt = Instant.now();
 %0A%0A      
@@ -4408,23 +4408,21 @@
 nt alive
-Seconds
+Count
  = 0;%0A  
@@ -4438,24 +4438,273 @@
 le (true) %7B%0A
+            Instant loopStartedAt = Instant.now();%0A            if (loopStartedAt.isAfter(terminateStartedAt.plus(Duration.ofMinutes(10)))) %7B%0A                throw new IllegalStateException(%22Server didn't shutdown within 10 minutes%22);%0A            %7D%0A%0A
             
@@ -4783,23 +4783,21 @@
    alive
-Seconds
+Count
 ++;%0A    
@@ -5332,187 +5332,226 @@
     
-if (aliveSeconds %3E Duration.ofMinutes(5).getSeconds()) %7B%0A                throw new Illegal
+// sleep for 1 second%0A            long sleepMillis = Duration.between(Instant.now(), loop
 Sta
+r
 te
-Exception(%22Server doesn't shutdown%22);%0A            %7D%0A%0A            Thread.sleep(1000);
+dAt.plusSeconds(1)).toMillis();%0A            if (sleepMillis %3E 0) %7B%0A                Thread.sleep(sleepMillis);%0A            %7D
 %0A   
@@ -5874,23 +5874,21 @@
 at(alive
-Seconds
+Count
 , greate
@@ -5893,17 +5893,17 @@
 terThan(
-5
+3
 ));%0A%0A   

Improve reliability of CI around ServerGracefulShutdownIT