/*
 * Copyright 2019 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.aws.s3;

import com.google.common.collect.ImmutableList;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.EnumSource;
import org.komamitsu.fluency.Fluency;
import org.komamitsu.fluency.aws.s3.ingester.AwsS3Ingester;
import org.komamitsu.fluency.aws.s3.ingester.DefaultS3DestinationDecider;
import org.komamitsu.fluency.aws.s3.ingester.sender.AwsS3Sender;
import org.komamitsu.fluency.aws.s3.recordformat.AwsS3RecordFormatter;
import org.komamitsu.fluency.aws.s3.recordformat.CsvRecordFormatter;
import org.komamitsu.fluency.aws.s3.recordformat.JsonlRecordFormatter;
import org.komamitsu.fluency.aws.s3.recordformat.MessagePackRecordFormatter;
import org.komamitsu.fluency.buffer.Buffer;
import org.komamitsu.fluency.flusher.Flusher;
import org.komamitsu.fluency.ingester.Ingester;
import org.komamitsu.fluency.recordformat.RecordFormatter;
import org.mockito.ArgumentCaptor;
import software.amazon.awssdk.regions.Region;

import java.time.ZoneId;
import java.time.ZoneOffset;

import static java.time.ZoneId.SHORT_IDS;
import static java.time.ZoneOffset.UTC;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;

class FluencyBuilderForAwsS3Test
{
    private FluencyBuilderForAwsS3 builderWithDefaultConfig;
    private FluencyBuilderForAwsS3 builderWithCustomConfig;
    private Fluency fluency;

    @BeforeEach
    void setUp()
    {
        fluency = mock(Fluency.class);

        builderWithDefaultConfig = spy(new FluencyBuilderForAwsS3());
        builderWithDefaultConfig.setFormatCsvColumnNames(ImmutableList.of("dummy"));
        doReturn(fluency).when(builderWithDefaultConfig)
                .createFluency(
                        any(RecordFormatter.class),
                        any(Ingester.class),
                        any(Buffer.Config.class),
                        any(Flusher.Config.class));

        {
            FluencyBuilderForAwsS3 builder = new FluencyBuilderForAwsS3();
            builder.setAwsEndpoint("https://foo.bar.org");
            builder.setAwsRegion("us-east-1");
            builder.setCompressionEnabled(false);
            builder.setAwsAccessKeyId("ACCESSKEYID");
            builder.setAwsSecretAccessKey("SECRETACCESSKEY");
            builder.setBufferChunkInitialSize(2 * 1024 * 1024);
            builder.setBufferChunkRetentionSize(9 * 1024 * 1024);
            builder.setBufferChunkRetentionTimeMillis(1234);
            builder.setMaxBufferSize(42 * 1024 * 1024L);
            builder.setS3KeyPrefix("mydata");
            builder.setS3KeySuffix(".zzz");
            builder.setS3KeyTimeZoneId(ZoneOffset.of("JST", SHORT_IDS));
            builder.setSenderRetryMax(4);
            builder.setSenderMaxRetryIntervalMillis(65432);
            builder.setSenderRetryIntervalMillis(543);
            builder.setSenderRetryFactor(1.234f);
            builder.setSenderWorkBufSize(99 * 1024);
            builderWithCustomConfig = spy(builder);
        }
        doReturn(fluency).when(builderWithCustomConfig)
                .createFluency(
                        any(RecordFormatter.class),
                        any(Ingester.class),
                        any(Buffer.Config.class),
                        any(Flusher.Config.class));
    }

    @ParameterizedTest
    @EnumSource(FluencyBuilderForAwsS3.FormatType.class)
    void buildWithDefaultConfig(FluencyBuilderForAwsS3.FormatType formatType)
    {
        FluencyBuilderForAwsS3 builder = builderWithDefaultConfig;
        AwsS3Sender sender = mock(AwsS3Sender.class);
        doReturn(sender).when(builder).createSender(any(AwsS3Sender.Config.class));

        builder.setFormatType(formatType);
        Fluency fluency = builder.build();
        assertEquals(fluency, this.fluency);

        ArgumentCaptor<AwsS3Sender.Config> configArgumentCaptor = ArgumentCaptor.forClass(AwsS3Sender.Config.class);
        verify(builder, times(1)).createSender(configArgumentCaptor.capture());
        AwsS3Sender.Config senderConfig = configArgumentCaptor.getValue();
        assertNull(senderConfig.getEndpoint());
        assertNull(senderConfig.getRegion());
        assertNull(senderConfig.getAwsAccessKeyId());
        assertNull(senderConfig.getAwsSecretAccessKey());
        assertEquals(10, senderConfig.getRetryMax());
        assertEquals(1000, senderConfig.getRetryIntervalMs());
        assertEquals(30000, senderConfig.getMaxRetryIntervalMs());
        assertEquals(2.0, senderConfig.getRetryFactor());
        assertEquals(8192, senderConfig.getWorkBufSize());
        assertTrue(senderConfig.isCompressionEnabled());

        ArgumentCaptor<RecordFormatter> recordFormatterArgumentCaptor = ArgumentCaptor.forClass(RecordFormatter.class);
        ArgumentCaptor<Ingester> ingesterArgumentCaptor = ArgumentCaptor.forClass(Ingester.class);
        verify(builder, times(1)).buildFromIngester(
                recordFormatterArgumentCaptor.capture(), ingesterArgumentCaptor.capture());

        RecordFormatter recordFormatter = recordFormatterArgumentCaptor.getValue();
        Class<? extends AwsS3RecordFormatter> expectedAwsS3RecordFormatter = null;
        String expectedS3KeySuffix = null;
        switch (formatType) {
            case MESSAGE_PACK:
                expectedAwsS3RecordFormatter = MessagePackRecordFormatter.class;
                expectedS3KeySuffix = ".msgpack.gz";
                break;
            case JSONL:
                expectedAwsS3RecordFormatter = JsonlRecordFormatter.class;
                expectedS3KeySuffix = ".jsonl.gz";
                break;
            case CSV:
                expectedAwsS3RecordFormatter = CsvRecordFormatter.class;
                expectedS3KeySuffix = ".csv.gz";
                break;
        }
        assertEquals(expectedAwsS3RecordFormatter, recordFormatter.getClass());

        AwsS3Ingester ingester = (AwsS3Ingester) ingesterArgumentCaptor.getValue();
        assertEquals(sender, ingester.getSender());
        DefaultS3DestinationDecider destinationDecider =
                (DefaultS3DestinationDecider) ingester.getS3DestinationDecider();
        assertNull(destinationDecider.getKeyPrefix());
        assertEquals(expectedS3KeySuffix, destinationDecider.getKeySuffix());
        assertEquals(UTC, destinationDecider.getZoneId());

        ArgumentCaptor<Buffer.Config> bufferConfigArgumentCaptor = ArgumentCaptor.forClass(Buffer.Config.class);
        ArgumentCaptor<Flusher.Config> flusherConfigArgumentCaptor = ArgumentCaptor.forClass(Flusher.Config.class);
        verify(builder, times(1)).createFluency(
                eq(recordFormatter),
                eq(ingester),
                bufferConfigArgumentCaptor.capture(),
                flusherConfigArgumentCaptor.capture());

        Buffer.Config bufferConfig = bufferConfigArgumentCaptor.getValue();
        assertEquals(512 * 1024 * 1024, bufferConfig.getMaxBufferSize());
        assertEquals(4 * 1024 * 1024, bufferConfig.getChunkInitialSize());
        assertEquals(64 * 1024 * 1024, bufferConfig.getChunkRetentionSize());
        assertEquals(30 * 1000, bufferConfig.getChunkRetentionTimeMillis());
    }

    @Test
    void buildWithCustomConfig()
    {
        FluencyBuilderForAwsS3 builder = builderWithCustomConfig;
        AwsS3Sender sender = mock(AwsS3Sender.class);
        doReturn(sender).when(builder).createSender(any(AwsS3Sender.Config.class));

        builder.setFormatType(FluencyBuilderForAwsS3.FormatType.JSONL);
        Fluency fluency = builder.build();
        assertEquals(fluency, this.fluency);

        ArgumentCaptor<AwsS3Sender.Config> configArgumentCaptor = ArgumentCaptor.forClass(AwsS3Sender.Config.class);
        verify(builder, times(1)).createSender(configArgumentCaptor.capture());
        AwsS3Sender.Config senderConfig = configArgumentCaptor.getValue();
        assertEquals("https://foo.bar.org", senderConfig.getEndpoint());
        assertEquals(Region.US_EAST_1.toString(), senderConfig.getRegion());
        assertEquals("ACCESSKEYID", senderConfig.getAwsAccessKeyId());
        assertEquals("SECRETACCESSKEY", senderConfig.getAwsSecretAccessKey());
        assertEquals(4, senderConfig.getRetryMax());
        assertEquals(543, senderConfig.getRetryIntervalMs());
        assertEquals(65432, senderConfig.getMaxRetryIntervalMs());
        assertEquals(1.234f, senderConfig.getRetryFactor());
        assertEquals(99 * 1024, senderConfig.getWorkBufSize());
        assertFalse(senderConfig.isCompressionEnabled());

        ArgumentCaptor<RecordFormatter> recordFormatterArgumentCaptor = ArgumentCaptor.forClass(RecordFormatter.class);
        ArgumentCaptor<Ingester> ingesterArgumentCaptor = ArgumentCaptor.forClass(Ingester.class);
        verify(builder, times(1)).buildFromIngester(
                recordFormatterArgumentCaptor.capture(), ingesterArgumentCaptor.capture());

        RecordFormatter recordFormatter = recordFormatterArgumentCaptor.getValue();
        assertTrue(recordFormatter instanceof JsonlRecordFormatter);

        AwsS3Ingester ingester = (AwsS3Ingester) ingesterArgumentCaptor.getValue();
        assertEquals(sender, ingester.getSender());
        DefaultS3DestinationDecider destinationDecider =
                (DefaultS3DestinationDecider) ingester.getS3DestinationDecider();
        assertEquals("mydata", destinationDecider.getKeyPrefix());
        assertEquals(".zzz", destinationDecider.getKeySuffix());
        assertEquals(ZoneId.of("JST", SHORT_IDS), destinationDecider.getZoneId());

        ArgumentCaptor<Buffer.Config> bufferConfigArgumentCaptor = ArgumentCaptor.forClass(Buffer.Config.class);
        ArgumentCaptor<Flusher.Config> flusherConfigArgumentCaptor = ArgumentCaptor.forClass(Flusher.Config.class);
        verify(builder, times(1)).createFluency(
                eq(recordFormatter),
                eq(ingester),
                bufferConfigArgumentCaptor.capture(),
                flusherConfigArgumentCaptor.capture());

        Buffer.Config bufferConfig = bufferConfigArgumentCaptor.getValue();
        assertEquals(42 * 1024 * 1024, bufferConfig.getMaxBufferSize());
        assertEquals(2 * 1024 * 1024, bufferConfig.getChunkInitialSize());
        assertEquals(9 * 1024 * 1024, bufferConfig.getChunkRetentionSize());
        assertEquals(1234, bufferConfig.getChunkRetentionTimeMillis());
    }
}
