/*
 * 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.metadata.etcd.helpers;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;

import io.etcd.jetcd.ByteSequence;
import io.etcd.jetcd.options.PutOption;
import io.etcd.jetcd.support.CloseableClient;
import io.etcd.jetcd.support.Observers;
import java.nio.charset.StandardCharsets;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.function.Consumer;
import java.util.function.Function;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.common.concurrent.FutureUtils;
import org.apache.bookkeeper.metadata.etcd.testing.EtcdTestBase;
import org.apache.bookkeeper.versioning.LongVersion;
import org.apache.bookkeeper.versioning.Version.Occurred;
import org.apache.bookkeeper.versioning.Versioned;
import org.apache.commons.compress.utils.Sets;
import org.apache.commons.lang3.RandomStringUtils;
import org.junit.Test;

/**
 * Integration test {@link KeySetReader}.
 */
@Slf4j
public class KeySetReaderTest extends EtcdTestBase {

    private static final Function<ByteSequence, String> BYTE_SEQUENCE_STRING_FUNCTION =
            bs -> bs.toString(StandardCharsets.UTF_8);

    @Test
    public void testReadSingleKey() throws Exception {
        String key = RandomStringUtils.randomAlphabetic(16);
        ByteSequence keyBs = ByteSequence.from(key, StandardCharsets.UTF_8);
        try (KeySetReader<String> ksReader = new KeySetReader<>(
            etcdClient,
            BYTE_SEQUENCE_STRING_FUNCTION,
            keyBs,
            null
        )) {
            // key not exists
            Versioned<Set<String>> versionedKeys = FutureUtils.result(ksReader.read());
            assertTrue(
                "VersionedKeys : " + versionedKeys,
                ((LongVersion) versionedKeys.getVersion()).getLongVersion() > 0L);
            assertEquals(0, versionedKeys.getValue().size());
            assertFalse(ksReader.isWatcherSet());

            // keys should be cached
            assertEquals(versionedKeys, ksReader.getLocalValue());

            // update a value
            String value = RandomStringUtils.randomAlphabetic(32);
            ByteSequence valueBs = ByteSequence.from(value, StandardCharsets.UTF_8);
            FutureUtils.result(etcdClient.getKVClient().put(keyBs, valueBs));

            // update the value should not change local value
            assertEquals(versionedKeys, ksReader.getLocalValue());

            // read the key again
            Versioned<Set<String>> newVersionedKey = FutureUtils.result(ksReader.read());
            assertEquals(Occurred.AFTER, newVersionedKey.getVersion().compare(versionedKeys.getVersion()));
            assertEquals(1, newVersionedKey.getValue().size());
            assertEquals(Sets.newHashSet(key), newVersionedKey.getValue());

            // local value should be changed
            assertEquals(newVersionedKey, ksReader.getLocalValue());
        }
    }

    @Test
    public void testWatchSingleKey() throws Exception {
        String key = RandomStringUtils.randomAlphabetic(16);
        ByteSequence keyBs = ByteSequence.from(key, StandardCharsets.UTF_8);
        KeySetReader<String> ksReader = null;
        try {
            ksReader = new KeySetReader<>(
                etcdClient,
                BYTE_SEQUENCE_STRING_FUNCTION,
                keyBs,
                null
            );
            LinkedBlockingQueue<Versioned<Set<String>>> notifications = new LinkedBlockingQueue<>();
            Consumer<Versioned<Set<String>>> keyConsumer = consumeVersionedKeySet(notifications);

            // key not exists
            Versioned<Set<String>> versionedKeys = FutureUtils.result(ksReader.readAndWatch(keyConsumer));
            assertTrue(
                "VersionedKeys : " + versionedKeys,
                ((LongVersion) versionedKeys.getVersion()).getLongVersion() > 0L);
            assertEquals(0, versionedKeys.getValue().size());
            assertTrue(ksReader.isWatcherSet());

            // keys should be cached
            assertEquals(versionedKeys, ksReader.getLocalValue());
            Versioned<Set<String>> newVersionedKey = notifications.take();
            assertEquals(Occurred.CONCURRENTLY, newVersionedKey.getVersion().compare(versionedKeys.getVersion()));
            assertEquals(versionedKeys, newVersionedKey);
            versionedKeys = newVersionedKey;

            // update a value
            String value = RandomStringUtils.randomAlphabetic(32);
            ByteSequence valueBs = ByteSequence.from(value, StandardCharsets.UTF_8);
            FutureUtils.result(etcdClient.getKVClient().put(keyBs, valueBs));

            // we should get notified with updated key set
            newVersionedKey = notifications.take();
            assertEquals(Occurred.AFTER, newVersionedKey.getVersion().compare(versionedKeys.getVersion()));
            assertEquals(1, newVersionedKey.getValue().size());
            assertEquals(Sets.newHashSet(key), newVersionedKey.getValue());

            // local value should be changed
            assertEquals(newVersionedKey, ksReader.getLocalValue());
            versionedKeys = newVersionedKey;

            // delete the key
            FutureUtils.result(etcdClient.getKVClient().delete(keyBs));
            newVersionedKey = notifications.take();
            assertEquals(Occurred.AFTER, newVersionedKey.getVersion().compare(versionedKeys.getVersion()));
            assertEquals(0, newVersionedKey.getValue().size());

            // local value should be changed
            assertEquals(newVersionedKey, ksReader.getLocalValue());
        } finally {
            if (null != ksReader) {
                ksReader.close();
            }
        }
        assertNotNull(ksReader);
        assertFalse(ksReader.isWatcherSet());
    }

    @Test
    public void testWatchSingleKeyWithTTL() throws Exception {
        String key = RandomStringUtils.randomAlphabetic(16);
        ByteSequence keyBs = ByteSequence.from(key, StandardCharsets.UTF_8);
        KeySetReader<String> ksReader = null;
        try {
            ksReader = new KeySetReader<>(
                etcdClient,
                BYTE_SEQUENCE_STRING_FUNCTION,
                keyBs,
                null
            );
            LinkedBlockingQueue<Versioned<Set<String>>> notifications = new LinkedBlockingQueue<>();
            Consumer<Versioned<Set<String>>> keyConsumer = consumeVersionedKeySet(notifications);

            // key not exists
            Versioned<Set<String>> versionedKeys = FutureUtils.result(ksReader.readAndWatch(keyConsumer));
            assertTrue(
                "VersionedKeys : " + versionedKeys,
                ((LongVersion) versionedKeys.getVersion()).getLongVersion() > 0L);
            assertEquals(0, versionedKeys.getValue().size());
            assertTrue(ksReader.isWatcherSet());

            // keys should be cached
            assertEquals(versionedKeys, ksReader.getLocalValue());
            // no watch event should be issued
            Versioned<Set<String>> newVersionedKey = notifications.take();
            assertEquals(Occurred.CONCURRENTLY, newVersionedKey.getVersion().compare(versionedKeys.getVersion()));
            assertEquals(versionedKeys, newVersionedKey);
            versionedKeys = newVersionedKey;

            // create a key with ttl
            long leaseId = FutureUtils.result(etcdClient.getLeaseClient().grant(1)).getID();
            String value = RandomStringUtils.randomAlphabetic(32);
            ByteSequence valueBs = ByteSequence.from(value, StandardCharsets.UTF_8);
            FutureUtils.result(etcdClient.getKVClient()
                .put(keyBs, valueBs, PutOption.newBuilder().withLeaseId(leaseId).build()));

            // we should get notified with updated key set
            newVersionedKey = notifications.take();
            assertEquals(Occurred.AFTER, newVersionedKey.getVersion().compare(versionedKeys.getVersion()));
            assertEquals(1, newVersionedKey.getValue().size());
            assertEquals(Sets.newHashSet(key), newVersionedKey.getValue());

            // local value should be changed
            assertEquals(newVersionedKey, ksReader.getLocalValue());
            versionedKeys = newVersionedKey;

            // the key will be deleted after TTL
            newVersionedKey = notifications.take();
            assertEquals(Occurred.AFTER, newVersionedKey.getVersion().compare(versionedKeys.getVersion()));
            assertEquals(0, newVersionedKey.getValue().size());

            // local value should be changed
            assertEquals(newVersionedKey, ksReader.getLocalValue());
        } finally {
            if (null != ksReader) {
                ksReader.close();
            }
        }
        assertNotNull(ksReader);
        assertFalse(ksReader.isWatcherSet());
    }

    @Test
    public void testReadKeySet() throws Exception {
        String prefix = RandomStringUtils.randomAlphabetic(16);
        ByteSequence beginKeyBs = ByteSequence.from(prefix + "-000", StandardCharsets.UTF_8);
        ByteSequence endKeyBs = ByteSequence.from(prefix + "-999", StandardCharsets.UTF_8);
        try (KeySetReader<String> ksReader = new KeySetReader<>(
            etcdClient,
            BYTE_SEQUENCE_STRING_FUNCTION,
            beginKeyBs,
            endKeyBs
        )) {
            // key not exists
            Versioned<Set<String>> versionedKeys = FutureUtils.result(ksReader.read());
            assertTrue(
                "VersionedKeys : " + versionedKeys,
                ((LongVersion) versionedKeys.getVersion()).getLongVersion() > 0L);
            assertEquals(0, versionedKeys.getValue().size());
            assertFalse(ksReader.isWatcherSet());

            // keys should be cached
            assertEquals(versionedKeys, ksReader.getLocalValue());

            Set<String> expectedKeySet = new HashSet<>();
            for (int i = 0; i < 20; i++) {
                // update a value
                String key = String.format("%s-%03d", prefix, i);
                String value = RandomStringUtils.randomAlphabetic(32);
                ByteSequence keyBs = ByteSequence.from(key, StandardCharsets.UTF_8);
                ByteSequence valueBs = ByteSequence.from(value, StandardCharsets.UTF_8);
                expectedKeySet.add(key);
                FutureUtils.result(etcdClient.getKVClient().put(keyBs, valueBs));

                // update the value should not change local value
                assertEquals(versionedKeys, ksReader.getLocalValue());

                // read the key again
                Versioned<Set<String>> newVersionedKey = FutureUtils.result(ksReader.read());
                assertEquals(Occurred.AFTER, newVersionedKey.getVersion().compare(versionedKeys.getVersion()));
                assertEquals(expectedKeySet, newVersionedKey.getValue());

                // local value should be changed
                assertEquals(newVersionedKey, ksReader.getLocalValue());
                versionedKeys = newVersionedKey;
            }
        }
    }

    @Test
    public void testWatchKeySet() throws Exception {
        String prefix = RandomStringUtils.randomAlphabetic(16);
        ByteSequence beginKeyBs = ByteSequence.from(prefix + "-000", StandardCharsets.UTF_8);
        ByteSequence endKeyBs = ByteSequence.from(prefix + "-999", StandardCharsets.UTF_8);
        KeySetReader<String> ksReader = null;
        try {
            ksReader = new KeySetReader<>(
                etcdClient,
                BYTE_SEQUENCE_STRING_FUNCTION,
                beginKeyBs,
                endKeyBs
            );
            LinkedBlockingQueue<Versioned<Set<String>>> notifications = new LinkedBlockingQueue<>();
            Consumer<Versioned<Set<String>>> keyConsumer = consumeVersionedKeySet(notifications);

            // key not exists
            Versioned<Set<String>> versionedKeys = FutureUtils.result(ksReader.readAndWatch(keyConsumer));
            assertTrue(
                "VersionedKeys : " + versionedKeys,
                ((LongVersion) versionedKeys.getVersion()).getLongVersion() > 0L);
            assertEquals(0, versionedKeys.getValue().size());
            assertTrue(ksReader.isWatcherSet());

            // keys should be cached
            assertEquals(versionedKeys, ksReader.getLocalValue());
            Versioned<Set<String>> newVersionedKey = notifications.take();
            assertEquals(Occurred.CONCURRENTLY, newVersionedKey.getVersion().compare(versionedKeys.getVersion()));
            assertEquals(versionedKeys, newVersionedKey);
            versionedKeys = newVersionedKey;

            Set<String> expectedKeySet = new HashSet<>();
            for (int i = 0; i < 20; i++) {
                // update a value
                String key = String.format("%s-%03d", prefix, i);
                String value = RandomStringUtils.randomAlphabetic(32);
                ByteSequence keyBs = ByteSequence.from(key, StandardCharsets.UTF_8);
                ByteSequence valueBs = ByteSequence.from(value, StandardCharsets.UTF_8);
                expectedKeySet.add(key);
                FutureUtils.result(etcdClient.getKVClient().put(keyBs, valueBs));

                // we should get notified with updated key set
                newVersionedKey = notifications.take();
                assertEquals(Occurred.AFTER, newVersionedKey.getVersion().compare(versionedKeys.getVersion()));
                assertEquals(expectedKeySet, newVersionedKey.getValue());

                // local value should be changed
                assertEquals(newVersionedKey, ksReader.getLocalValue());
                versionedKeys = newVersionedKey;
            }

            for (int i = 0; i < 20; i++) {
                // delete the key
                String key = String.format("%s-%03d", prefix, i);
                ByteSequence keyBs = ByteSequence.from(key, StandardCharsets.UTF_8);
                expectedKeySet.remove(key);
                FutureUtils.result(etcdClient.getKVClient().delete(keyBs));

                // we should get notified with updated key set
                newVersionedKey = notifications.take();
                assertEquals(Occurred.AFTER, newVersionedKey.getVersion().compare(versionedKeys.getVersion()));
                assertEquals(expectedKeySet, newVersionedKey.getValue());

                // local value should be changed
                assertEquals(newVersionedKey, ksReader.getLocalValue());
                versionedKeys = newVersionedKey;
            }
        } finally {
            if (null != ksReader) {
                ksReader.close();
            }
        }
        assertNotNull(ksReader);
        assertFalse(ksReader.isWatcherSet());
    }

    @Test
    public void testWatchKeySetWithTTL() throws Exception {
        String prefix = RandomStringUtils.randomAlphabetic(16);
        ByteSequence beginKeyBs = ByteSequence.from(prefix + "-000", StandardCharsets.UTF_8);
        ByteSequence endKeyBs = ByteSequence.from(prefix + "-999", StandardCharsets.UTF_8);
        KeySetReader<String> ksReader = null;
        try {
            ksReader = new KeySetReader<>(
                etcdClient,
                BYTE_SEQUENCE_STRING_FUNCTION,
                beginKeyBs,
                endKeyBs
            );
            LinkedBlockingQueue<Versioned<Set<String>>> notifications = new LinkedBlockingQueue<>();
            Consumer<Versioned<Set<String>>> keyConsumer = consumeVersionedKeySet(notifications);

            // key not exists
            Versioned<Set<String>> versionedKeys = FutureUtils.result(ksReader.readAndWatch(keyConsumer));
            assertTrue(
                "VersionedKeys : " + versionedKeys,
                ((LongVersion) versionedKeys.getVersion()).getLongVersion() > 0L);
            assertEquals(0, versionedKeys.getValue().size());
            assertTrue(ksReader.isWatcherSet());

            // keys should be cached
            assertEquals(versionedKeys, ksReader.getLocalValue());
            // no watch event should be issued
            Versioned<Set<String>> newVersionedKey = notifications.take();
            assertEquals(Occurred.CONCURRENTLY, newVersionedKey.getVersion().compare(versionedKeys.getVersion()));
            assertEquals(versionedKeys, newVersionedKey);
            versionedKeys = newVersionedKey;

            // create keys with ttl
            long leaseId = FutureUtils.result(etcdClient.getLeaseClient().grant(1)).getID();
            CloseableClient ka = etcdClient.getLeaseClient().keepAlive(leaseId, Observers.observer(response -> {
            }));

            Set<String> expectedKeySet = new HashSet<>();
            for (int i = 0; i < 20; i++) {
                String key = String.format("%s-%03d", prefix, i);
                String value = RandomStringUtils.randomAlphabetic(32);
                ByteSequence keyBs = ByteSequence.from(key, StandardCharsets.UTF_8);
                ByteSequence valueBs = ByteSequence.from(value, StandardCharsets.UTF_8);
                expectedKeySet.add(key);
                FutureUtils.result(etcdClient.getKVClient()
                    .put(keyBs, valueBs, PutOption.newBuilder().withLeaseId(leaseId).build()));

                // we should get notified with updated key set
                newVersionedKey = notifications.take();
                assertEquals(Occurred.AFTER, newVersionedKey.getVersion().compare(versionedKeys.getVersion()));
                assertEquals(expectedKeySet, newVersionedKey.getValue());

                // local value should be changed
                assertEquals(newVersionedKey, ksReader.getLocalValue());
                versionedKeys = newVersionedKey;
            }

            // stop keep alive all the keys should be expired.
            ka.close();

            // all the keys will be deleted after TTL in same batch.
            newVersionedKey = notifications.take();
            // local value should be changed
            assertEquals(newVersionedKey, ksReader.getLocalValue());
            assertEquals(Occurred.AFTER, newVersionedKey.getVersion().compare(versionedKeys.getVersion()));
            assertTrue(newVersionedKey.getValue().isEmpty());
        } finally {
            if (null != ksReader) {
                ksReader.close();
            }
        }
        assertNotNull(ksReader);
        assertFalse(ksReader.isWatcherSet());
    }
}
