/*
 * Copyright 2018 Mitsunori Komatsu (komamitsu)
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.komamitsu.fluency.fluentd;

import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.databind.SerializerProvider;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.fasterxml.jackson.databind.ser.std.StdSerializer;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.MethodSource;
import org.komamitsu.fluency.Fluency;
import org.komamitsu.fluency.FluencyBuilder;
import org.komamitsu.fluency.fluentd.ingester.sender.RetryableSender;
import org.komamitsu.fluency.fluentd.recordformat.FluentdRecordFormatter;
import org.komamitsu.fluency.ingester.Ingester;
import org.msgpack.core.MessagePack;
import org.msgpack.core.MessagePacker;
import org.msgpack.core.MessageUnpacker;
import org.msgpack.value.ImmutableRawValue;
import org.msgpack.value.StringValue;
import org.msgpack.value.Value;
import org.msgpack.value.ValueFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Stream;

import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.instanceOf;
import static org.hamcrest.Matchers.is;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.komamitsu.fluency.fluentd.SSLTestSocketFactories.SSL_CLIENT_SOCKET_FACTORY;
import static org.mockito.Mockito.mock;

class FluencyTest
{
    private static final Logger LOG = LoggerFactory.getLogger(FluencyTest.class);
    private static final StringValue KEY_OPTION_SIZE = ValueFactory.newString("size");
    private static final StringValue KEY_OPTION_CHUNK = ValueFactory.newString("chunk");
    private Ingester ingester;

    @BeforeEach
    void setUp()
    {
        ingester = mock(Ingester.class);
    }

    @Test
    void testSenderErrorHandler()
            throws IOException, InterruptedException
    {
        final CountDownLatch countDownLatch = new CountDownLatch(1);
        final AtomicReference<Throwable> errorContainer = new AtomicReference<>();

        FluencyBuilderForFluentd builder = new FluencyBuilderForFluentd();
        builder.setSenderMaxRetryCount(1);
        builder.setErrorHandler(e -> {
            errorContainer.set(e);
            countDownLatch.countDown();
        });

        try (Fluency fluency = builder.build(Integer.MAX_VALUE)) {
            HashMap<String, Object> event = new HashMap<>();
            event.put("name", "foo");
            fluency.emit("tag", event);

            if (!countDownLatch.await(10, TimeUnit.SECONDS)) {
                throw new AssertionError("Timeout");
            }

            assertThat(errorContainer.get(), is(instanceOf(RetryableSender.RetryOverException.class)));
        }
    }

    static Stream<Boolean> sslFlagsProvider()
    {
        return Stream.of(false, true);
    }

    @ParameterizedTest
    @MethodSource("sslFlagsProvider")
    void testWithoutAckResponse(final boolean sslEnabled)
            throws Throwable
    {
        Exception exception = new ConfigurableTestServer(sslEnabled).run(
                clientSocket -> {
                    MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(clientSocket.getInputStream());
                    assertEquals(3, unpacker.unpackArrayHeader());
                    assertEquals("foo.bar", unpacker.unpackString());
                    ImmutableRawValue rawValue = unpacker.unpackValue().asRawValue();
                    Map<Value, Value> map = unpacker.unpackValue().asMapValue().map();
                    assertEquals(1, map.size());
                    assertEquals(rawValue.asByteArray().length, map.get(KEY_OPTION_SIZE).asIntegerValue().asInt());
                    unpacker.close();
                },
                serverPort -> {
                    FluencyBuilderForFluentd builder = new FluencyBuilderForFluentd();
                    builder.setSslEnabled(sslEnabled);
                    if (sslEnabled) {
                        builder.setSslSocketFactory(SSL_CLIENT_SOCKET_FACTORY);
                    }

                    try (Fluency fluency = builder.build(serverPort)) {
                        fluency.emit("foo.bar", new HashMap<>());
                    }
                }, 5000);
        assertNull(exception);
    }

    @ParameterizedTest
    @MethodSource("sslFlagsProvider")
    void testWithAckResponseButNotReceiveToken(final boolean sslEnabled)
            throws Throwable
    {
        Exception exception = new ConfigurableTestServer(sslEnabled).run(
                clientSocket -> {
                    MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(clientSocket.getInputStream());
                    assertEquals(3, unpacker.unpackArrayHeader());
                    assertEquals("foo.bar", unpacker.unpackString());
                    ImmutableRawValue rawValue = unpacker.unpackValue().asRawValue();
                    Map<Value, Value> map = unpacker.unpackValue().asMapValue().map();
                    assertEquals(2, map.size());
                    assertEquals(rawValue.asByteArray().length, map.get(KEY_OPTION_SIZE).asIntegerValue().asInt());
                    assertNotNull(map.get(KEY_OPTION_CHUNK).asRawValue().asString());
                    unpacker.close();
                },
                serverPort -> {
                    FluencyBuilderForFluentd builder = new FluencyBuilderForFluentd();
                    builder.setSslEnabled(sslEnabled);
                    builder.setAckResponseMode(true);
                    if (sslEnabled) {
                        builder.setSslSocketFactory(SSL_CLIENT_SOCKET_FACTORY);
                    }

                    try (Fluency fluency = builder.build(serverPort)) {
                        fluency.emit("foo.bar", new HashMap<>());
                    }
                }, 5000);
        assertEquals(exception.getClass(), TimeoutException.class);
    }

    @ParameterizedTest
    @MethodSource("sslFlagsProvider")
    void testWithAckResponseButWrongReceiveToken(final boolean sslEnabled)
            throws Throwable
    {
        Exception exception = new ConfigurableTestServer(sslEnabled).run(
                clientSocket -> {
                    MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(clientSocket.getInputStream());
                    assertEquals(3, unpacker.unpackArrayHeader());
                    assertEquals("foo.bar", unpacker.unpackString());
                    ImmutableRawValue rawValue = unpacker.unpackValue().asRawValue();
                    Map<Value, Value> map = unpacker.unpackValue().asMapValue().map();
                    assertEquals(2, map.size());
                    assertEquals(rawValue.asByteArray().length, map.get(KEY_OPTION_SIZE).asIntegerValue().asInt());
                    assertNotNull(map.get(KEY_OPTION_CHUNK).asRawValue().asString());

                    MessagePacker packer = MessagePack.newDefaultPacker(clientSocket.getOutputStream());
                    packer.packMapHeader(1)
                            .packString("ack").packString(UUID.randomUUID().toString())
                            .close();

                    // Close the input stream after closing the output stream to avoid closing a socket too early
                    unpacker.close();
                },
                serverPort -> {
                    FluencyBuilderForFluentd builder = new FluencyBuilderForFluentd();
                    builder.setSslEnabled(sslEnabled);
                    builder.setAckResponseMode(true);
                    if (sslEnabled) {
                        builder.setSslSocketFactory(SSL_CLIENT_SOCKET_FACTORY);
                    }

                    try (Fluency fluency = builder.build(serverPort)) {
                        fluency.emit("foo.bar", new HashMap<>());
                    }
                }, 5000);
        assertEquals(exception.getClass(), TimeoutException.class);
    }

    @ParameterizedTest
    @MethodSource("sslFlagsProvider")
    public void testWithAckResponseWithProperToken(final boolean sslEnabled)
            throws Throwable
    {
        Exception exception = new ConfigurableTestServer(sslEnabled).run(
                clientSocket -> {
                    MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(clientSocket.getInputStream());
                    assertEquals(3, unpacker.unpackArrayHeader());
                    assertEquals("foo.bar", unpacker.unpackString());
                    ImmutableRawValue rawValue = unpacker.unpackValue().asRawValue();
                    Map<Value, Value> map = unpacker.unpackValue().asMapValue().map();
                    assertEquals(2, map.size());
                    assertEquals(rawValue.asByteArray().length, map.get(KEY_OPTION_SIZE).asIntegerValue().asInt());
                    String ackResponseToken = map.get(KEY_OPTION_CHUNK).asRawValue().asString();
                    assertNotNull(ackResponseToken);

                    MessagePacker packer = MessagePack.newDefaultPacker(clientSocket.getOutputStream());
                    packer.packMapHeader(1)
                            .packString("ack").packString(ackResponseToken)
                            .close();

                    // Close the input stream after closing the output stream to avoid closing a socket too early
                    unpacker.close();
                },
                serverPort -> {
                    FluencyBuilderForFluentd builder = new FluencyBuilderForFluentd();
                    builder.setSslEnabled(sslEnabled);
                    builder.setAckResponseMode(true);
                    if (sslEnabled) {
                        builder.setSslSocketFactory(SSL_CLIENT_SOCKET_FACTORY);
                    }

                    try (Fluency fluency = builder.build(serverPort)) {
                        fluency.emit("foo.bar", new HashMap<>());
                    }
                }, 5000);
        assertNull(exception);
    }

    @ParameterizedTest
    @MethodSource("sslFlagsProvider")
    void testBufferWithJacksonModule()
            throws IOException
    {
        AtomicBoolean serialized = new AtomicBoolean();
        SimpleModule simpleModule = new SimpleModule();
        simpleModule.addSerializer(Foo.class, new FooSerializer(serialized));

        FluentdRecordFormatter.Config recordFormatterConfig = new FluentdRecordFormatter.Config();
        recordFormatterConfig.setJacksonModules(Collections.singletonList(simpleModule));

        Fluency fluency = new FluencyBuilder()
                .buildFromIngester(new FluentdRecordFormatter(recordFormatterConfig), ingester);

        Map<String, Object> event = new HashMap<>();
        Foo foo = new Foo();
        foo.s = "Hello";
        event.put("foo", foo);
        fluency.emit("tag", event);

        assertThat(serialized.get(), is(true));
    }

    static class Foo
    {
        String s;
    }

    public static class FooSerializer
            extends StdSerializer<Foo>
    {
        final AtomicBoolean serialized;

        FooSerializer(AtomicBoolean serialized)
        {
            super(Foo.class);
            this.serialized = serialized;
        }

        @Override
        public void serialize(Foo value, JsonGenerator gen, SerializerProvider provider)
                throws IOException
        {
            gen.writeStartObject();
            gen.writeStringField("s", "Foo:" + value.s);
            gen.writeEndObject();
            serialized.set(true);
        }
    }
}
