/*
 * 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.bookkeeper.statelib.impl.kv;

import static java.nio.charset.StandardCharsets.UTF_8;
import static org.apache.bookkeeper.common.concurrent.FutureUtils.result;
import static org.apache.bookkeeper.statelib.testing.executors.MockExecutorController.THREAD_NAME_PREFIX;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import static org.mockito.Mockito.mock;

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import java.io.File;
import java.net.URI;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.time.Duration;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.common.coder.ByteArrayCoder;
import org.apache.bookkeeper.common.concurrent.FutureUtils;
import org.apache.bookkeeper.statelib.api.StateStoreSpec;
import org.apache.bookkeeper.statelib.api.exceptions.InvalidStateStoreException;
import org.apache.bookkeeper.statelib.impl.rocksdb.RocksUtils;
import org.apache.bookkeeper.statelib.impl.rocksdb.checkpoint.fs.FSCheckpointManager;
import org.apache.bookkeeper.statelib.testing.executors.MockExecutorController;
import org.apache.commons.io.FileUtils;
import org.apache.distributedlog.DLMTestUtil;
import org.apache.distributedlog.DLSN;
import org.apache.distributedlog.TestDistributedLogBase;
import org.apache.distributedlog.api.namespace.Namespace;
import org.apache.distributedlog.api.namespace.NamespaceBuilder;
import org.junit.After;
import org.junit.AfterClass;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.junit.rules.TestName;

/**
 * Unit test of {@link RocksdbKVStore}.
 */
@Slf4j
public class TestRocksdbKVAsyncStoreWithCheckpoints extends TestDistributedLogBase {

    private static URI uri;
    private static Namespace namespace;

    @BeforeClass
    public static void setupCluster() throws Exception {
        TestDistributedLogBase.setupCluster();
        uri = DLMTestUtil.createDLMURI(zkPort, "/mvcc");
        conf.setPeriodicFlushFrequencyMilliSeconds(2);
        conf.setWriteLockEnabled(false);
        conf.setExplicitTruncationByApplication(true);
        namespace = NamespaceBuilder.newBuilder()
            .conf(conf)
            .uri(uri)
            .clientId("test-mvcc-async-store")
            .build();
    }

    @AfterClass
    public static void teardownCluster() throws Exception {
        if (null != namespace) {
            namespace.close();
        }
        TestDistributedLogBase.teardownCluster();
    }

    @Rule
    public final TestName runtime = new TestName();
    @Rule
    public final TemporaryFolder testDir = new TemporaryFolder();

    private String streamName;
    private RocksdbKVAsyncStore<byte[], byte[]> store;
    private ScheduledExecutorService mockWriteExecutor;
    private ExecutorService realWriteExecutor;
    private MockExecutorController writeExecutorController;
    private ScheduledExecutorService checkpointExecutor;
    private FSCheckpointManager checkpointStore;
    private File localDir;
    private File remoteDir;

    @Before
    @Override
    public void setup() throws Exception {
        super.setup();
        ensureURICreated(uri);

        localDir = testDir.newFolder("local");
        remoteDir = testDir.newFolder("remote");

        checkpointStore = new FSCheckpointManager(remoteDir);

        // initialize the scheduler
        realWriteExecutor = Executors.newSingleThreadExecutor(
                new ThreadFactoryBuilder().setNameFormat(THREAD_NAME_PREFIX + "%d").build());
        mockWriteExecutor = mock(ScheduledExecutorService.class);
        writeExecutorController = new MockExecutorController(realWriteExecutor)
            .controlExecute(mockWriteExecutor)
            .controlSubmit(mockWriteExecutor)
            .controlSchedule(mockWriteExecutor)
            .controlScheduleAtFixedRate(mockWriteExecutor, 5);

        checkpointExecutor = Executors.newSingleThreadScheduledExecutor();

        store = new RocksdbKVAsyncStore<>(
            () -> new RocksdbKVStore<>(),
            () -> namespace);
        FutureUtils.result(store.init(initSpec(runtime.getMethodName())));
    }

    private StateStoreSpec initSpec(String streamName) {
        return StateStoreSpec.builder()
            .name(streamName)
            .keyCoder(ByteArrayCoder.of())
            .valCoder(ByteArrayCoder.of())
            .stream(streamName)
            .localStateStoreDir(localDir)
            .writeIOScheduler(mockWriteExecutor)
            .checkpointStore(checkpointStore)
            .checkpointIOScheduler(checkpointExecutor)
            .checkpointDuration(Duration.ofMinutes(1))
            .checkpointChecksumEnable(true)
            .checkpointChecksumCompatible(false)
            .build();
    }

    @After
    @Override
    public void teardown() throws Exception {
        if (null != streamName) {
            namespace.deleteLog(streamName);
        }

        if (null != store) {
            store.close();
        }

        if (null != checkpointExecutor) {
            checkpointExecutor.shutdown();
        }

        if (null != realWriteExecutor) {
            realWriteExecutor.shutdown();
        }

        if (null != checkpointStore) {
            checkpointStore.close();
        }

        super.teardown();
    }

    @Test(expected = NullPointerException.class)
    public void testInitMissingStreamName() throws Exception {
        this.streamName = "test-init-missing-stream-name";
        StateStoreSpec spec = StateStoreSpec.builder()
            .name(streamName)
            .keyCoder(ByteArrayCoder.of())
            .valCoder(ByteArrayCoder.of())
            .localStateStoreDir(localDir)
            .build();
        result(store.init(spec));
    }

    private byte[] getKey(int i) {
        return String.format("key-%05d", i).getBytes(UTF_8);
    }

    private byte[] getValue(int i) {
        return String.format("value-%05d", i).getBytes(UTF_8);
    }

    private void writeNumKvs(int numKvs, int startIdx) throws Exception {
        CompletableFuture<Void> lastWrite = null;
        for (int i = 0; i < numKvs; i++) {
            byte[] key = getKey(startIdx + i);
            byte[] val = getValue(startIdx + i);

            lastWrite = store.put(key, val);
        }
        FutureUtils.result(lastWrite);
    }

    private void readNumKvs(int numKvs, int startIdx) throws Exception {
        for (int i = 0; i < numKvs; i++) {
            byte[] key = getKey(startIdx + i);

            byte[] expectedVal = getValue(startIdx + i);
            byte[] actualVal = FutureUtils.result(store.get(key));

            assertArrayEquals(expectedVal, actualVal);
        }
    }


    @Test
    public void testRestartNoCheckpoint() throws Exception {
        final int numKvs = 100;

        // write 100 kvs
        writeNumKvs(numKvs, 0);

        // close the store
        store.close();

        // since the lock isn't advanced so no checkpoint is created.
        List<String> files = checkpointStore.listFiles(RocksUtils.getDestCheckpointsPath(store.name()));
        assertTrue(files.isEmpty());

        // remove local dir
        FileUtils.deleteDirectory(new File(localDir.getAbsolutePath()));

        // reload the store
        store = new RocksdbKVAsyncStore<>(
            () -> new RocksdbKVStore<>(),
            () -> namespace);
        FutureUtils.result(store.init(initSpec(runtime.getMethodName())));

        assertTrue(Files.exists(
            Paths.get(localDir.getAbsolutePath())));

        // verify the 100 kvs
        readNumKvs(numKvs, 0);
    }

    @Test
    public void testRestartWithCheckpoint() throws Exception {
        final int numKvs = 100;

        // write 100 kvs
        writeNumKvs(numKvs, 0);

        // advance the clock to trigger checkpointing
        writeExecutorController.advance(Duration.ofMinutes(1));

        // ensure checkpoint completed
        checkpointExecutor.submit(() -> {
        }).get();

        store.close();

        List<String> files = checkpointStore.listFiles(RocksUtils.getDestCheckpointsPath(store.name()));
        assertEquals(1, files.size());

        // remove local dir
        FileUtils.deleteDirectory(new File(localDir.getAbsolutePath()));

        // reload the store
        store = new RocksdbKVAsyncStore<>(
            () -> new RocksdbKVStore<>(),
            () -> namespace);
        FutureUtils.result(store.init(initSpec(runtime.getMethodName())));

        assertTrue(Files.exists(
            Paths.get(localDir.getAbsolutePath())));

        // verify the 100 kvs
        readNumKvs(numKvs, 0);
    }

    private void reinitStore(StateStoreSpec spec) throws Exception {
        store.close();

        // remove local dir
        FileUtils.deleteDirectory(new File(localDir.getAbsolutePath()));

        store = new RocksdbKVAsyncStore<>(
            () -> new RocksdbKVStore<>(),
            () -> namespace);
        result(store.init(spec));
    }

    @Test
    public void testMissingData() throws Exception {
        int numKeys = 100;
        this.streamName = runtime.getMethodName();
        StateStoreSpec spec = initSpec(runtime.getMethodName());

        reinitStore(spec);

        writeNumKvs(numKeys, 0);
        readNumKvs(numKeys, 0);

        // Reload store from Journal
        reinitStore(spec);

        readNumKvs(numKeys, 0);

        // Trigger a checkpoint
        store.getLocalStore().checkpoint();
        // ensure checkpoint completed
        checkpointExecutor.submit(() -> {}).get();

        writeNumKvs(numKeys, 100);

        reinitStore(spec);

        // Validate that we can load from both checkpoint and journal
        readNumKvs(numKeys, 0);
        readNumKvs(numKeys, 100);

        // Add some more data
        writeNumKvs(numKeys, 200);
        DLSN dlsn = store.getLastDLSN();

        // Truncate and purge part of Journal
        result(store.truncateJournal(dlsn));
        store.purgeOlderThan(Integer.MAX_VALUE);

        // reloading the statestore should fail.
        try {
            reinitStore(spec);
            fail("Store initialization should fail due to missing data");
        } catch (InvalidStateStoreException isse) {
            assertEquals(
                "replayJournal failed: Invalid starting transaction 204 expecting 102 for stream testMissingData",
                isse.getMessage());
        }
    }
}
