/*
 * 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.buffer;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.komamitsu.fluency.JsonRecordFormatter;
import org.komamitsu.fluency.ingester.Ingester;
import org.komamitsu.fluency.recordformat.RecordFormatter;
import org.komamitsu.fluency.util.Tuple;
import org.mockito.ArgumentCaptor;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;

import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.number.OrderingComparison.greaterThan;
import static org.hamcrest.number.OrderingComparison.lessThan;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
import static org.mockito.ArgumentMatchers.*;
import static org.mockito.Mockito.*;

class BufferTest
{
    private static final Logger LOG = LoggerFactory.getLogger(BufferTest.class);
    private Ingester ingester;
    private RecordFormatter recordFormatter;
    private Buffer.Config bufferConfig;

    @BeforeEach
    void setUp()
    {
        recordFormatter = new JsonRecordFormatter();
        ingester = mock(Ingester.class);
        bufferConfig = new Buffer.Config();
    }

    @Test
    void testBuffer()
            throws IOException
    {
        int recordSize = 20; // '{"name":"komamitsu"}'
        int tags = 10;
        float allocRatio = 2f;
        int maxAllocSize = (recordSize * 4) * tags;

        bufferConfig.setChunkInitialSize(recordSize);
        bufferConfig.setChunkExpandRatio(allocRatio);
        bufferConfig.setMaxBufferSize(maxAllocSize);
        bufferConfig.setChunkRetentionSize(maxAllocSize / 2);
        Buffer buffer = new Buffer(bufferConfig, recordFormatter);

        assertEquals(0, buffer.getAllocatedSize());
        assertEquals(0, buffer.getBufferedDataSize());
        assertEquals(0f, buffer.getBufferUsage(), 0.001);

        Map<String, Object> data = new HashMap<>();
        data.put("name", "komamitsu");
        for (int i = 0; i < tags; i++) {
            buffer.append("foodb.bartbl" + i, 1420070400, data);
        }
        assertEquals(recordSize * tags, buffer.getAllocatedSize());
        assertEquals(recordSize * tags, buffer.getBufferedDataSize());
        assertEquals(((float) recordSize) * tags / maxAllocSize, buffer.getBufferUsage(), 0.001);
        for (int i = 0; i < tags; i++) {
            buffer.append("foodb.bartbl" + i, 1420070400, data);
        }
        float totalAllocatedSize = recordSize * (1 + allocRatio) * tags;
        assertEquals(totalAllocatedSize, buffer.getAllocatedSize(), 0.001);
        assertEquals(2L * recordSize * tags, buffer.getBufferedDataSize());
        assertEquals(totalAllocatedSize / maxAllocSize, buffer.getBufferUsage(), 0.001);

        buffer.flush(ingester, true);
        assertEquals(totalAllocatedSize, buffer.getAllocatedSize(), 0.001);
        assertEquals(0, buffer.getBufferedDataSize());
        assertEquals(totalAllocatedSize / maxAllocSize, buffer.getBufferUsage(), 0.001);

        buffer.close();
        assertEquals(0, buffer.getAllocatedSize());
        assertEquals(0, buffer.getBufferedDataSize());
        assertEquals(0, buffer.getBufferUsage(), 0.001);
    }

    @Test
    void testStrictBufferRetentionSizeManagement()
            throws IOException
    {
        int recordSize = 20; // '{"name":"komamitsu"}'
        int initSize = 15;
        float allocRatio = 2f;
        // 15 -> 30 -> 60 -> 120 -> 240
        int maxAllocSize = (int) (initSize * allocRatio * allocRatio * allocRatio * allocRatio);
        String tag = "foodb.bartbl";

        int retentionSize = 60;

        bufferConfig.setChunkInitialSize(initSize);
        bufferConfig.setChunkExpandRatio(allocRatio);
        bufferConfig.setMaxBufferSize(maxAllocSize);
        bufferConfig.setChunkRetentionSize(retentionSize);
        Buffer buffer = new Buffer(bufferConfig, recordFormatter);

        Map<String, Object> data = new HashMap<>();
        data.put("name", "komamitsu");

        // Buffered data size is 0 and additional 20 bytes data will be buffered (total size: 20 bytes).
        buffer.append(tag, 1420070400, data);
        buffer.flush(ingester, false);
        verify(ingester, times(0)).ingest(anyString(), any(ByteBuffer.class));
        assertEquals(20, buffer.getBufferedDataSize());

        // Buffered data size is 20 and additional 20 bytes data will be buffered (total size: 40 bytes).
        buffer.append(tag, 1420070400, data);
        buffer.flush(ingester, false);
        verify(ingester, times(0)).ingest(anyString(), any(ByteBuffer.class));
        assertEquals(40, buffer.getBufferedDataSize());

        // Buffered data size is 40 and additional 20 bytes data will be buffered (total size: 60 bytes).
        buffer.append(tag, 1420070400, data);
        buffer.flush(ingester, false);
        verify(ingester, times(0)).ingest(anyString(), any(ByteBuffer.class));
        assertEquals(60, buffer.getBufferedDataSize());

        // Buffered data size is 60 and additional 20 bytes exceeds the retention size.
        // The buffered 60 bytes data will be flushed and the additional data will be buffered (total size: 20 bytes).
        doAnswer(invocation -> {
            String receivedTag = invocation.getArgument(0);
            ByteBuffer receivedBuffer = invocation.getArgument(1);
            assertEquals(tag, receivedTag);
            assertEquals(60, receivedBuffer.remaining());
            return null;
        }).when(ingester).ingest(anyString(), any(ByteBuffer.class));

        buffer.append(tag, 1420070400, data);
        buffer.flush(ingester, false);
        verify(ingester, times(1)).ingest(anyString(), any(ByteBuffer.class));
        assertEquals(20, buffer.getBufferedDataSize());
    }

    @Test
    void flush()
            throws IOException
    {
        Buffer buffer = new Buffer(bufferConfig, recordFormatter);

        Ingester ingester = mock(Ingester.class);
        buffer.flush(ingester, false);

        verify(ingester, times(0)).ingest(anyString(), any(ByteBuffer.class));

        Map<String, Object> data = new HashMap<>();
        data.put("name", "komamitsu");
        buffer.append("foodb.bartbl", 1420070400, data);

        buffer.flush(ingester, false);

        verify(ingester, times(0)).ingest(anyString(), any(ByteBuffer.class));

        buffer.flush(ingester, true);

        verify(ingester, times(1)).ingest(eq("foodb.bartbl"), any(ByteBuffer.class));
    }

    @Test
    void testFileBackup()
            throws IOException
    {
        bufferConfig.setFileBackupDir(System.getProperty("java.io.tmpdir"));
        bufferConfig.setFileBackupPrefix("FileBackupTest");

        // Just for cleaning backup files
        try (Buffer buffer = new Buffer(bufferConfig, recordFormatter)) {
            buffer.clearBackupFiles();
        }
        try (Buffer buffer = new Buffer(bufferConfig, recordFormatter)) {
            assertEquals(0, buffer.getBufferedDataSize());
        }

        long currentTime = System.currentTimeMillis() / 1000;
        Map<String, Object> event0 = ImmutableMap.of("name", "a", "age", 42);
        Map<String, Object> event1 = ImmutableMap.of("name", "b", "age", 99);
        try (Buffer buffer = new Buffer(bufferConfig, recordFormatter)) {
            buffer.append("foo", currentTime, event0);
            buffer.append("bar", currentTime, event1);
        }

        Ingester ingester = mock(Ingester.class);
        try (Buffer buffer = new Buffer(bufferConfig, recordFormatter)) {
            buffer.flushInternal(ingester, true);
        }

        ObjectMapper objectMapper = new ObjectMapper();
        for (Tuple<String, Map<String, Object>> tagAndEvent :
                ImmutableList.of(new Tuple<>("foo", event0), new Tuple<>("bar", event1))) {
            ArgumentCaptor<ByteBuffer> byteBufferArgumentCaptor = ArgumentCaptor.forClass(ByteBuffer.class);
            verify(ingester, times(1)).ingest(eq(tagAndEvent.getFirst()), byteBufferArgumentCaptor.capture());
            ByteBuffer byteBuffer = byteBufferArgumentCaptor.getValue();
            byte[] bytes = new byte[byteBuffer.remaining()];
            byteBuffer.get(bytes);
            Map<String, Object> map = objectMapper.readValue(bytes, new TypeReference<Map<String, Object>>() {});
            assertEquals(tagAndEvent.getSecond(), map);
        }
    }

    @Test
    void testFileBackupThatIsNotDirectory()
            throws IOException
    {
        File backupDirFile = File.createTempFile("testFileBackupWithInvalidDir", ".tmp");
        backupDirFile.deleteOnExit();

        bufferConfig.setFileBackupDir(backupDirFile.getAbsolutePath());

        try (Buffer ignored = new Buffer(bufferConfig, recordFormatter)) {
            fail();
        }
        catch (IllegalArgumentException e) {
            assertTrue(true);
        }
    }

    @Test
    void testFileBackupThatIsWritable()
    {
        File backupDir = new File(System.getProperties().getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
        if (backupDir.mkdir()) {
            LOG.info("Created directory: {}", backupDir);
        }
        if (!backupDir.setWritable(false)) {
            throw new RuntimeException("Failed to revoke writable permission");
        }
        backupDir.deleteOnExit();

        bufferConfig.setFileBackupDir(backupDir.getAbsolutePath());

        try (Buffer ignored = new Buffer(bufferConfig, recordFormatter)) {
            fail();
        }
        catch (IllegalArgumentException e) {
            assertTrue(true);
        }
    }

    @Test
    void testFileBackupThatIsReadable()
    {
        File backupDir = new File(System.getProperties().getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
        if (backupDir.mkdir()) {
            LOG.info("Created directory: {}", backupDir);
        }
        if (!backupDir.setReadable(false)) {
            throw new RuntimeException("Failed to revoke readable permission");
        }
        backupDir.deleteOnExit();

        bufferConfig.setFileBackupDir(backupDir.getAbsolutePath());

        try (Buffer ignored = new Buffer(bufferConfig, recordFormatter)) {
            fail();
        }
        catch (IllegalArgumentException e) {
            assertTrue(true);
        }
    }

    @Test
    void testGetAllocatedSize()
            throws IOException
    {
        bufferConfig.setChunkInitialSize(256 * 1024);
        try (Buffer buffer = new Buffer(bufferConfig, recordFormatter)) {
            assertThat(buffer.getAllocatedSize(), is(0L));
            Map<String, Object> map = new HashMap<>();
            map.put("name", "komamitsu");
            for (int i = 0; i < 10; i++) {
                buffer.append("foo.bar", new Date().getTime(), map);
            }
            assertThat(buffer.getAllocatedSize(), is(256 * 1024L));
        }
    }

    @Test
    void testGetBufferedDataSize()
            throws IOException
    {
        bufferConfig.setChunkInitialSize(256 * 1024);
        try (Buffer buffer = new Buffer(bufferConfig, recordFormatter)) {
            assertThat(buffer.getBufferedDataSize(), is(0L));

            Map<String, Object> map = new HashMap<>();
            map.put("name", "komamitsu");
            for (int i = 0; i < 10; i++) {
                buffer.append("foo.bar", new Date().getTime(), map);
            }
            assertThat(buffer.getBufferedDataSize(), is(greaterThan(0L)));
            assertThat(buffer.getBufferedDataSize(), is(lessThan(512L)));

            buffer.flush(ingester, true);
            assertThat(buffer.getBufferedDataSize(), is(0L));
        }
    }

    @Test
    void testAppendIfItDoesNotThrowBufferOverflow()
            throws IOException
    {
        bufferConfig.setChunkInitialSize(64 * 1024);
        try (Buffer buffer = new Buffer(bufferConfig, recordFormatter)) {

            StringBuilder buf = new StringBuilder();

            for (int i = 0; i < 1024 * 60; i++) {
                buf.append('x');
            }
            String str60kb = buf.toString();

            for (int i = 0; i < 1024 * 40; i++) {
                buf.append('x');
            }
            String str100kb = buf.toString();

            {
                Map<String, Object> map = new HashMap<>();
                map.put("k", str60kb);
                buffer.append("tag0", new Date().getTime(), map);
            }

            {
                Map<String, Object> map = new HashMap<>();
                map.put("k", str100kb);
                buffer.append("tag0", new Date().getTime(), map);
            }
        }
    }

    @Test
    void validateConfig()
    {
        {
            Buffer.Config config = new Buffer.Config();
            config.setChunkInitialSize(4 * 1024 * 1024);
            config.setChunkRetentionSize(4 * 1024 * 1024);
            assertThrows(IllegalArgumentException.class, () -> new Buffer(config, recordFormatter));
        }

        {
            Buffer.Config config = new Buffer.Config();
            config.setChunkRetentionSize(64 * 1024 * 1024);
            config.setMaxBufferSize(64 * 1024 * 1024);
            assertThrows(IllegalArgumentException.class, () -> new Buffer(config, recordFormatter));
        }

        {
            Buffer.Config config = new Buffer.Config();
            config.setChunkExpandRatio(1.19f);
            assertThrows(IllegalArgumentException.class, () -> new Buffer(config, recordFormatter));
        }

        {
            Buffer.Config config = new Buffer.Config();
            config.setChunkRetentionTimeMillis(49);
            assertThrows(IllegalArgumentException.class, () -> new Buffer(config, recordFormatter));
        }
    }
}
