/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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.apache.kafka.connect.storage;

import org.apache.kafka.connect.errors.ConnectException;
import org.apache.kafka.connect.util.Callback;
import org.easymock.Capture;
import org.easymock.EasyMock;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.powermock.api.easymock.PowerMock;
import org.powermock.api.easymock.annotation.Mock;
import org.powermock.modules.junit4.PowerMockRunner;

import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;

@RunWith(PowerMockRunner.class)
public class OffsetStorageWriterTest {
    private static final String NAMESPACE = "namespace";
    // Connect format - any types should be accepted here
    private static final Map<String, String> OFFSET_KEY = Collections.singletonMap("key", "key");
    private static final Map<String, Integer> OFFSET_VALUE = Collections.singletonMap("key", 12);

    // Serialized
    private static final byte[] OFFSET_KEY_SERIALIZED = "key-serialized".getBytes();
    private static final byte[] OFFSET_VALUE_SERIALIZED = "value-serialized".getBytes();

    @Mock private OffsetBackingStore store;
    @Mock private Converter keyConverter;
    @Mock private Converter valueConverter;
    private OffsetStorageWriter writer;

    private static Exception exception = new RuntimeException("error");

    private ExecutorService service;

    @Before
    public void setup() {
        writer = new OffsetStorageWriter(store, NAMESPACE, keyConverter, valueConverter);
        service = Executors.newFixedThreadPool(1);
    }

    @After
    public void teardown() {
        service.shutdownNow();
    }

    @Test
    public void testWriteFlush() throws Exception {
        @SuppressWarnings("unchecked")
        Callback<Void> callback = PowerMock.createMock(Callback.class);
        expectStore(OFFSET_KEY, OFFSET_KEY_SERIALIZED, OFFSET_VALUE, OFFSET_VALUE_SERIALIZED, callback, false, null);

        PowerMock.replayAll();

        writer.offset(OFFSET_KEY, OFFSET_VALUE);

        assertTrue(writer.beginFlush());
        writer.doFlush(callback).get(1000, TimeUnit.MILLISECONDS);

        PowerMock.verifyAll();
    }

    // It should be possible to set offset values to null
    @Test
    public void testWriteNullValueFlush() throws Exception {
        @SuppressWarnings("unchecked")
        Callback<Void> callback = PowerMock.createMock(Callback.class);
        expectStore(OFFSET_KEY, OFFSET_KEY_SERIALIZED, null, null, callback, false, null);

        PowerMock.replayAll();

        writer.offset(OFFSET_KEY, null);

        assertTrue(writer.beginFlush());
        writer.doFlush(callback).get(1000, TimeUnit.MILLISECONDS);

        PowerMock.verifyAll();
    }

    // It should be possible to use null keys. These aren't actually stored as null since the key is wrapped to include
    // info about the namespace (connector)
    @Test
    public void testWriteNullKeyFlush() throws Exception {
        @SuppressWarnings("unchecked")
        Callback<Void> callback = PowerMock.createMock(Callback.class);
        expectStore(null, null, OFFSET_VALUE, OFFSET_VALUE_SERIALIZED, callback, false, null);

        PowerMock.replayAll();

        writer.offset(null, OFFSET_VALUE);

        assertTrue(writer.beginFlush());
        writer.doFlush(callback).get(1000, TimeUnit.MILLISECONDS);

        PowerMock.verifyAll();
    }

    @Test
    public void testNoOffsetsToFlush() {
        // If no offsets are flushed, we should finish immediately and not have made any calls to the
        // underlying storage layer

        PowerMock.replayAll();

        // Should not return a future
        assertFalse(writer.beginFlush());

        PowerMock.verifyAll();
    }

    @Test
    public void testFlushFailureReplacesOffsets() throws Exception {
        // When a flush fails, we shouldn't just lose the offsets. Instead, they should be restored
        // such that a subsequent flush will write them.

        @SuppressWarnings("unchecked")
        final Callback<Void> callback = PowerMock.createMock(Callback.class);
        // First time the write fails
        expectStore(OFFSET_KEY, OFFSET_KEY_SERIALIZED, OFFSET_VALUE, OFFSET_VALUE_SERIALIZED, callback, true, null);
        // Second time it succeeds
        expectStore(OFFSET_KEY, OFFSET_KEY_SERIALIZED, OFFSET_VALUE, OFFSET_VALUE_SERIALIZED, callback, false, null);
        // Third time it has no data to flush so we won't get past beginFlush()

        PowerMock.replayAll();

        writer.offset(OFFSET_KEY, OFFSET_VALUE);
        assertTrue(writer.beginFlush());
        writer.doFlush(callback).get(1000, TimeUnit.MILLISECONDS);
        assertTrue(writer.beginFlush());
        writer.doFlush(callback).get(1000, TimeUnit.MILLISECONDS);
        assertFalse(writer.beginFlush());

        PowerMock.verifyAll();
    }

    @Test
    public void testAlreadyFlushing() {
        @SuppressWarnings("unchecked")
        final Callback<Void> callback = PowerMock.createMock(Callback.class);
        // Trigger the send, but don't invoke the callback so we'll still be mid-flush
        CountDownLatch allowStoreCompleteCountdown = new CountDownLatch(1);
        expectStore(OFFSET_KEY, OFFSET_KEY_SERIALIZED, OFFSET_VALUE, OFFSET_VALUE_SERIALIZED, null, false, allowStoreCompleteCountdown);

        PowerMock.replayAll();

        writer.offset(OFFSET_KEY, OFFSET_VALUE);
        assertTrue(writer.beginFlush());
        writer.doFlush(callback);
        assertThrows(ConnectException.class, writer::beginFlush);

        PowerMock.verifyAll();
    }

    @Test
    public void testCancelBeforeAwaitFlush() {
        PowerMock.replayAll();

        writer.offset(OFFSET_KEY, OFFSET_VALUE);
        assertTrue(writer.beginFlush());
        writer.cancelFlush();

        PowerMock.verifyAll();
    }

    @Test
    public void testCancelAfterAwaitFlush() throws Exception {
        @SuppressWarnings("unchecked")
        Callback<Void> callback = PowerMock.createMock(Callback.class);
        CountDownLatch allowStoreCompleteCountdown = new CountDownLatch(1);
        // In this test, the write should be cancelled so the callback will not be invoked and is not
        // passed to the expectStore call
        expectStore(OFFSET_KEY, OFFSET_KEY_SERIALIZED, OFFSET_VALUE, OFFSET_VALUE_SERIALIZED, null, false, allowStoreCompleteCountdown);

        PowerMock.replayAll();

        writer.offset(OFFSET_KEY, OFFSET_VALUE);
        assertTrue(writer.beginFlush());
        // Start the flush, then immediately cancel before allowing the mocked store request to finish
        Future<Void> flushFuture = writer.doFlush(callback);
        writer.cancelFlush();
        allowStoreCompleteCountdown.countDown();
        flushFuture.get(1000, TimeUnit.MILLISECONDS);

        PowerMock.verifyAll();
    }

    /**
     * Expect a request to store data to the underlying OffsetBackingStore.
     *
     * @param key the key for the offset
     * @param keySerialized serialized version of the key
     * @param value the value for the offset
     * @param valueSerialized serialized version of the value
     * @param callback the callback to invoke when completed, or null if the callback isn't
     *                 expected to be invoked
     * @param fail if true, treat
     * @param waitForCompletion if non-null, a CountDownLatch that should be awaited on before
     *                          invoking the callback. A (generous) timeout is still imposed to
     *                          ensure tests complete.
     * @return the captured set of ByteBuffer key-value pairs passed to the storage layer
     */
    private void expectStore(Map<String, String> key, byte[] keySerialized,
                             Map<String, Integer> value, byte[] valueSerialized,
                             final Callback<Void> callback,
                             final boolean fail,
                             final CountDownLatch waitForCompletion) {
        List<Object> keyWrapped = Arrays.asList(NAMESPACE, key);
        EasyMock.expect(keyConverter.fromConnectData(NAMESPACE, null, keyWrapped)).andReturn(keySerialized);
        EasyMock.expect(valueConverter.fromConnectData(NAMESPACE, null, value)).andReturn(valueSerialized);

        final Capture<Callback<Void>> storeCallback = Capture.newInstance();
        final Map<ByteBuffer, ByteBuffer> offsetsSerialized = Collections.singletonMap(
                keySerialized == null ? null : ByteBuffer.wrap(keySerialized),
                valueSerialized == null ? null : ByteBuffer.wrap(valueSerialized));
        EasyMock.expect(store.set(EasyMock.eq(offsetsSerialized), EasyMock.capture(storeCallback)))
            .andAnswer(() ->
                service.submit(() -> {
                    if (waitForCompletion != null)
                        assertTrue(waitForCompletion.await(10000, TimeUnit.MILLISECONDS));

                    if (fail) {
                        storeCallback.getValue().onCompletion(exception, null);
                    } else {
                        storeCallback.getValue().onCompletion(null, null);
                    }
                    return null;
                })
            );
        if (callback != null) {
            if (fail) {
                callback.onCompletion(EasyMock.eq(exception), EasyMock.eq(null));
            } else {
                callback.onCompletion(null, null);
            }
        }
        PowerMock.expectLastCall();
    }

}
