package io.digdag.core.agent;

import com.google.common.collect.ImmutableMap;
import com.google.inject.Inject;
import com.google.inject.Scopes;
import com.google.inject.multibindings.Multibinder;
import io.digdag.client.config.Config;
import io.digdag.client.config.ConfigUtils;
import io.digdag.core.DigdagEmbed;
import io.digdag.spi.Operator;
import io.digdag.spi.OperatorContext;
import io.digdag.spi.OperatorFactory;
import io.digdag.spi.TaskExecutionException;
import io.digdag.spi.TaskResult;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import java.nio.file.Paths;
import java.util.Map;

import static io.digdag.client.config.ConfigUtils.newConfig;
import static io.digdag.core.workflow.OperatorTestingUtils.newTaskRequest;
import static io.digdag.core.workflow.WorkflowTestingUtils.setupEmbed;
import static org.hamcrest.Matchers.instanceOf;
import static org.hamcrest.Matchers.is;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.fail;

public class OperatorManagerExceptionTest
{
    static class CustomErrorOperatorFactory
            implements OperatorFactory
    {
        @Inject
        public CustomErrorOperatorFactory()
        { }

        @Override
        public String getType()
        {
            return "custom_error";
        }

        @Override
        public Operator newOperator(OperatorContext context)
        {
            return new CustomErrorOperator(context);
        }
    }

    static class CustomErrorOperator
            implements Operator
    {
        private final OperatorContext context;

        public CustomErrorOperator(OperatorContext context)
        {
            this.context = context;
        }

        @Override
        public TaskResult run()
        {
            throw EXCEPTIONS.get(context.getTaskRequest().getConfig().get("_command", String.class));
        }
    }

    private static class CustomRuntimeException extends RuntimeException
    {
        CustomRuntimeException(String message)
        {
            super(message);
        }

        CustomRuntimeException(String message, Throwable cause)
        {
            super(message, cause);
        }
    }

    private final static Map<String, RuntimeException> EXCEPTIONS = ImmutableMap.<String, RuntimeException>builder()
        .put("runtime", new RuntimeException("foobar"))
        .put("custom", new CustomRuntimeException("foobar"))
        .put("custom_nested", new CustomRuntimeException(null, new RuntimeException("nested")))
        .put("wrapped_runtime", new TaskExecutionException(new RuntimeException("foobar")))
        .put("wrapped_custom", new TaskExecutionException(new CustomRuntimeException("foobar")))
        .put("wrapped_null_message", new TaskExecutionException(new CustomRuntimeException(null, new RuntimeException("cause"))))
        .put("wrapped_custom_message", new TaskExecutionException("custom!!", new RuntimeException("foo")))
        .build();

    private DigdagEmbed embed;
    private OperatorManager operatorManager;

    @Before
    public void setUp()
            throws Exception
    {
        this.embed = setupEmbed((bootstrap) -> bootstrap.addModules((binder) -> {
                    Multibinder.newSetBinder(binder, OperatorFactory.class)
                        .addBinding().to(CustomErrorOperatorFactory.class).in(Scopes.SINGLETON);
                })
            );
        this.operatorManager = embed.getInjector().getInstance(OperatorManager.class);
    }

    @After
    public void shutdown()
            throws Exception
    {
        embed.close();
    }

    @Test
    public void verifyException()
    {
        expectRuntimeException("runtime", RuntimeException.class, "foobar (runtime)");
        expectRuntimeException("custom", CustomRuntimeException.class, "foobar (custom runtime)");
        expectRuntimeException("custom_nested", CustomRuntimeException.class, "CustomRuntimeException (custom runtime)\n> nested (runtime)");
        expectExecutionException("wrapped_runtime", "foobar (runtime)");
        expectExecutionException("wrapped_custom", "foobar (custom runtime)");
        expectExecutionException("wrapped_null_message", "cause (custom runtime)");
        expectExecutionException("wrapped_custom_message", "custom!!");
    }

    private void expectExecutionException(String name, String expectedMessage)
    {
        Config config = newConfig().set("_command", name);
        try {
            operatorManager.callExecutor(Paths.get(""), "custom_error", newTaskRequest().withConfig(config));
            fail("expected TaskExecutionException");
        }
        catch (RuntimeException ex) {
            assertThat(ex, instanceOf(TaskExecutionException.class));
            Config error = ((TaskExecutionException) ex).getError(ConfigUtils.configFactory).or(newConfig());
            assertThat(error.get("message", String.class, null), is(expectedMessage));
        }
    }

    private void expectRuntimeException(String name, Class<?> expectedClass, String expectedMessage)
    {
        Config config = newConfig().set("_command", name);
        try {
            operatorManager.callExecutor(Paths.get(""), "custom_error", newTaskRequest().withConfig(config));
            fail("expected RuntimeException");
        }
        catch (RuntimeException ex) {
            assertThat((Object) ex.getClass(), is((Object) expectedClass));
            assertThat(OperatorManager.formatExceptionMessage(ex), is(expectedMessage));
        }
    }
}

@@ -303,24 +303,77 @@
 igdagEmbed;%0A
+import io.digdag.core.workflow.WorkflowTestingUtils;%0A
 import io.di
@@ -586,22 +586,81 @@
 g.junit.
-Before
+AfterClass;%0Aimport org.junit.Before;%0Aimport org.junit.BeforeClass
 ;%0Aimport
@@ -868,79 +868,8 @@
 st;%0A
-import static io.digdag.core.workflow.WorkflowTestingUtils.setupEmbed;%0A
 impo
@@ -3012,16 +3012,23 @@
 private 
+static 
 DigdagEm
@@ -3038,16 +3038,17 @@
  embed;%0A
+%0A
     priv
@@ -3096,16 +3096,21 @@
  @Before
+Class
 %0A    pub
@@ -3117,18 +3117,28 @@
 lic 
-void setUp
+static void shutdown
 ()%0A 
@@ -3183,21 +3183,16 @@
     
-this.
 embed = 
 setu
@@ -3187,16 +3187,37 @@
 embed = 
+WorkflowTestingUtils.
 setupEmb
@@ -3369,16 +3369,20 @@
         
+    
 .addBind
@@ -3477,99 +3477,133 @@
     
+);%0A
     
-);
+%7D%0A
 %0A    
-    this.operatorManager = embed.getInjector().getInstance(OperatorManager.class
+@AfterClass%0A    public static void destroyDigdagEmbed()%0A            throws Exception%0A    %7B%0A        embed.close(
 );%0A 
@@ -3613,21 +3613,22 @@
 %7D%0A%0A    @
-After
+Before
 %0A    pub
@@ -3629,39 +3629,36 @@
    public void s
-hutdown
+etUp
 ()%0A            t
@@ -3683,36 +3683,100 @@
   %7B%0A        
-embed.close(
+this.operatorManager = embed.getInjector().getInstance(OperatorManager.class
 );%0A    %7D%0A%0A  

Fix OperatorManagerExceptionTest not to recreate DigdagEmbed instance for each test method to make tests faster