package org.embulk.spi;

import static org.junit.Assert.assertEquals;
import static org.msgpack.value.ValueFactory.newBoolean;
import static org.msgpack.value.ValueFactory.newInteger;
import static org.msgpack.value.ValueFactory.newMap;
import static org.msgpack.value.ValueFactory.newString;

import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import org.embulk.config.ConfigDiff;
import org.embulk.config.ConfigSource;
import org.embulk.config.Task;
import org.embulk.config.TaskReport;
import org.embulk.config.TaskSource;
import org.embulk.test.EmbulkTestRuntime;
import org.embulk.test.MockFormatterPlugin;
import org.embulk.test.MockParserPlugin;
import org.embulk.test.PageTestUtils;
import org.junit.Rule;
import org.junit.Test;
import org.msgpack.value.ImmutableMapValue;

public class TestFileOutputRunner {
    @Rule
    public EmbulkTestRuntime runtime = new EmbulkTestRuntime();

    public interface PluginTask extends Task {}

    private static class MockFileOutputPlugin implements FileOutputPlugin {
        Boolean transactionCompleted = null;

        @Override
        public ConfigDiff transaction(ConfigSource config, int taskCount,
                FileOutputPlugin.Control control) {
            PluginTask task = config.loadConfig(PluginTask.class);
            control.run(task.dump());
            return Exec.newConfigDiff();
        }

        @Override
        public ConfigDiff resume(TaskSource taskSource,
                int taskCount,
                FileOutputPlugin.Control control) {
            throw new UnsupportedOperationException();
        }

        @Override
        public void cleanup(TaskSource taskSource,
                int taskCount,
                List<TaskReport> successTaskReports) {}

        @Override
        public TransactionalFileOutput open(TaskSource taskSource,
                final int taskIndex) {
            return new TransactionalFileOutput() {

                @Override
                public void nextFile() {}

                @Override
                public void add(Buffer buffer) {}

                @Override
                public void finish() {}

                @Override
                public void close() {}

                @Override
                public void abort() {
                    transactionCompleted = false;
                }

                @Override
                public TaskReport commit() {
                    transactionCompleted = true;
                    return Exec.newTaskReport();
                }
            };
        }
    }

    @Test
    public void testMockFormatterIteration() {
        MockFileOutputPlugin fileOutputPlugin = new MockFileOutputPlugin();
        final FileOutputRunner runner = new FileOutputRunner(fileOutputPlugin);

        final List<?> columns = TestUtils.listOf(
                TestUtils.mapOf("name", "col1", "type", "boolean", "option", TestUtils.mapOf()),
                TestUtils.mapOf("name", "col2", "type", "long", "option", TestUtils.mapOf()),
                TestUtils.mapOf("name", "col3", "type", "double", "option", TestUtils.mapOf()),
                TestUtils.mapOf("name", "col4", "type", "string", "option", TestUtils.mapOf()),
                TestUtils.mapOf("name", "col5", "type", "timestamp", "option", TestUtils.mapOf()),
                TestUtils.mapOf("name", "col6", "type", "json", "option", TestUtils.mapOf()));
        ConfigSource config = Exec
                .newConfigSource()
                .set("type", "unused?")
                .set("formatter",
                        TestUtils.mapOf("type", "mock", "columns", columns));
        final Schema schema = config.getNested("formatter")
                .loadConfig(MockParserPlugin.PluginTask.class)
                .getSchemaConfig().toSchema();

        runner.transaction(config, schema, 1, new OutputPlugin.Control() {
            public List<TaskReport> run(final TaskSource outputTask) {
                TransactionalPageOutput tran = runner.open(outputTask, schema,
                        1);
                boolean committed = false;
                try {
                    ImmutableMapValue jsonValue = newMap(
                            newString("_c1"), newBoolean(true),
                            newString("_c2"), newInteger(10),
                            newString("_c3"), newString("embulk"),
                            newString("_c4"), newMap(newString("k"), newString("v"))
                    );
                    for (Page page : PageTestUtils.buildPage(
                            runtime.getBufferAllocator(), schema, true, 2L,
                            3.0D, "45", Instant.ofEpochMilli(678L), jsonValue, true, 2L,
                            3.0D, "45", Instant.ofEpochMilli(678L), jsonValue)) {
                        tran.add(page);
                    }
                    tran.commit();
                    committed = true;
                } finally {
                    if (!committed) {
                        tran.abort();
                    }
                    tran.close();
                }
                return new ArrayList<TaskReport>();
            }
        });

        assertEquals(true, fileOutputPlugin.transactionCompleted);
        assertEquals(2, MockFormatterPlugin.records.size());
        for (List<Object> record : MockFormatterPlugin.records) {
            assertEquals(Boolean.TRUE, record.get(0));
            assertEquals(2L, record.get(1));
            assertEquals(3.0D, (Double) record.get(2), 0.1D);
            assertEquals("45", record.get(3));
            assertEquals(678L, ((Instant) record.get(4)).toEpochMilli());
            assertEquals("{\"_c1\":true,\"_c2\":10,\"_c3\":\"embulk\",\"_c4\":{\"k\":\"v\"}}", record.get(5).toString());
        }
    }

    @Test
    public void testTransactionAborted() {
        MockFileOutputPlugin fileOutputPlugin = new MockFileOutputPlugin();
        final FileOutputRunner runner = new FileOutputRunner(fileOutputPlugin);

        final List<?> columns = TestUtils.listOf(
                TestUtils.mapOf("name", "col1", "type", "boolean", "option", TestUtils.mapOf()),
                TestUtils.mapOf("name", "col2", "type", "long", "option", TestUtils.mapOf()),
                TestUtils.mapOf("name", "col3", "type", "double", "option", TestUtils.mapOf()),
                TestUtils.mapOf("name", "col4", "type", "string", "option", TestUtils.mapOf()),
                TestUtils.mapOf("name", "col5", "type", "timestamp", "option", TestUtils.mapOf()),
                TestUtils.mapOf("name", "col6", "type", "json", "option", TestUtils.mapOf()));
        ConfigSource config = Exec
                .newConfigSource()
                .set("type", "unused?")
                .set("formatter",
                        TestUtils.mapOf("type", "mock", "columns", columns));
        final Schema schema = config.getNested("formatter")
                .loadConfig(MockParserPlugin.PluginTask.class)
                .getSchemaConfig().toSchema();

        try {
            runner.transaction(config, schema, 1, new OutputPlugin.Control() {
                public List<TaskReport> run(final TaskSource outputTask) {
                    TransactionalPageOutput tran = runner.open(outputTask,
                            schema, 1);
                    boolean committed = false;
                    try {
                        tran.add(null);
                        tran.commit();
                        committed = true;
                    } finally {
                        if (!committed) {
                            tran.abort();
                        }
                        tran.close();
                    }
                    return new ArrayList<TaskReport>();
                }
            });
        } catch (NullPointerException npe) {
            // Just passing through.
        }

        assertEquals(false, fileOutputPlugin.transactionCompleted);
    }
}
