/**
 * 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.pulsar.common.naming;

import static org.mockito.Mockito.any;
import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertTrue;
import static org.testng.Assert.fail;

import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.SortedSet;
import java.util.concurrent.CompletableFuture;

import org.apache.commons.lang3.tuple.Pair;
import org.apache.pulsar.broker.PulsarService;
import org.apache.pulsar.broker.cache.LocalZooKeeperCacheService;
import org.apache.pulsar.common.naming.TopicName;
import org.apache.pulsar.common.naming.NamespaceBundle;
import org.apache.pulsar.common.naming.NamespaceBundleFactory;
import org.apache.pulsar.common.naming.NamespaceBundles;
import org.apache.pulsar.common.naming.NamespaceName;
import org.apache.pulsar.common.policies.data.LocalPolicies;
import org.apache.pulsar.zookeeper.ZooKeeperDataCache;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.annotations.Test;

import com.github.benmanes.caffeine.cache.AsyncLoadingCache;
import com.google.common.collect.BoundType;
import com.google.common.collect.Range;
import com.google.common.collect.Sets;
import com.google.common.hash.Hashing;

public class NamespaceBundlesTest {

    private final NamespaceBundleFactory factory = getNamespaceBundleFactory();

    @SuppressWarnings("unchecked")
    @Test
    public void testConstructor() throws Exception {
        try {
            new NamespaceBundles(null, (SortedSet<Long>) null, null);
            fail("Should fail w/ null pointer exception");
        } catch (NullPointerException npe) {
            // OK, expected
        }

        try {
            new NamespaceBundles(NamespaceName.get("pulsar/use/ns2"), (SortedSet<Long>) null, null);
            fail("Should fail w/ null pointer exception");
        } catch (NullPointerException npe) {
            // OK, expected
        }

        try {
            new NamespaceBundles(NamespaceName.get("pulsar.use.ns2"), (SortedSet<Long>) null, null);
            fail("Should fail w/ illegal argument exception");
        } catch (IllegalArgumentException iae) {
            // OK, expected
        }

        try {
            new NamespaceBundles(NamespaceName.get("pulsar/use/ns2"), (SortedSet<Long>) null, factory);
            fail("Should fail w/ null pointer exception");
        } catch (NullPointerException npe) {
            // OK, expected
        }

        SortedSet<Long> partitions = Sets.newTreeSet();
        try {
            new NamespaceBundles(NamespaceName.get("pulsar/use/ns2"), partitions, factory);
            fail("Should fail w/ illegal argument exception");
        } catch (IllegalArgumentException iae) {
            // OK, expected
        }

        partitions.add(0l);
        partitions.add(0x10000000l);
        partitions.add(0x40000000l);
        partitions.add(0xffffffffl);
        NamespaceBundles bundles = new NamespaceBundles(NamespaceName.get("pulsar/use/ns2"), partitions, factory);
        Field partitionField = NamespaceBundles.class.getDeclaredField("partitions");
        Field nsField = NamespaceBundles.class.getDeclaredField("nsname");
        Field bundlesField = NamespaceBundles.class.getDeclaredField("bundles");
        partitionField.setAccessible(true);
        nsField.setAccessible(true);
        bundlesField.setAccessible(true);
        long[] partFld = (long[]) partitionField.get(bundles);
        // the same instance
        assertEquals(partitions.size(), partFld.length);
        NamespaceName nsFld = (NamespaceName) nsField.get(bundles);
        assertEquals(nsFld.toString(), "pulsar/use/ns2");
        ArrayList<NamespaceBundle> bundleList = (ArrayList<NamespaceBundle>) bundlesField.get(bundles);
        assertEquals(bundleList.size(), 3);
        assertEquals(bundleList.get(0),
                factory.getBundle(nsFld, Range.range(0l, BoundType.CLOSED, 0x10000000l, BoundType.OPEN)));
        assertEquals(bundleList.get(1),
                factory.getBundle(nsFld, Range.range(0x10000000l, BoundType.CLOSED, 0x40000000l, BoundType.OPEN)));
        assertEquals(bundleList.get(2),
                factory.getBundle(nsFld, Range.range(0x40000000l, BoundType.CLOSED, 0xffffffffl, BoundType.CLOSED)));
    }

    @SuppressWarnings("unchecked")
    private NamespaceBundleFactory getNamespaceBundleFactory() {
        PulsarService pulsar = mock(PulsarService.class);
        LocalZooKeeperCacheService localZkCache = mock(LocalZooKeeperCacheService.class);
        ZooKeeperDataCache<LocalPolicies> poilciesCache = mock(ZooKeeperDataCache.class);
        when(pulsar.getLocalZkCacheService()).thenReturn(localZkCache);
        when(localZkCache.policiesCache()).thenReturn(poilciesCache);
        doNothing().when(poilciesCache).registerListener(any());
        return NamespaceBundleFactory.createFactory(pulsar, Hashing.crc32());
    }

    @Test
    public void testFindBundle() throws Exception {
        SortedSet<Long> partitions = Sets.newTreeSet();
        partitions.add(0l);
        partitions.add(0x40000000l);
        partitions.add(0xa0000000l);
        partitions.add(0xb0000000l);
        partitions.add(0xc0000000l);
        partitions.add(0xffffffffl);
        NamespaceBundles bundles = new NamespaceBundles(NamespaceName.get("pulsar/global/ns1"), partitions, factory);
        TopicName topicName = TopicName.get("persistent://pulsar/global/ns1/topic-1");
        NamespaceBundle bundle = bundles.findBundle(topicName);
        assertTrue(bundle.includes(topicName));

        topicName = TopicName.get("persistent://pulsar/use/ns2/topic-2");
        try {
            bundles.findBundle(topicName);
            fail("Should have failed due to mismatched namespace name");
        } catch (IllegalArgumentException iae) {
            // OK, expected
        }

        Long hashKey = factory.getLongHashCode(topicName.toString());
        // The following code guarantees that we have at least two ranges after the hashKey till the end
        SortedSet<Long> tailSet = partitions.tailSet(hashKey);
        tailSet.add(hashKey);
        // Now, remove the first range to ensure the hashKey is not included in <code>newPar</code>
        Iterator<Long> iter = tailSet.iterator();
        iter.next();
        SortedSet<Long> newPar = tailSet.tailSet(iter.next());

        try {
            bundles = new NamespaceBundles(topicName.getNamespaceObject(), newPar, factory);
            bundles.findBundle(topicName);
            fail("Should have failed due to out-of-range");
        } catch (IndexOutOfBoundsException iae) {
            // OK, expected
        }
    }

    @Test
    public void testsplitBundles() throws Exception {
        NamespaceName nsname = NamespaceName.get("pulsar/global/ns1");
        TopicName topicName = TopicName.get("persistent://pulsar/global/ns1/topic-1");
        NamespaceBundles bundles = factory.getBundles(nsname);
        NamespaceBundle bundle = bundles.findBundle(topicName);
        final int numberSplitBundles = 4;
        // (1) split in 4
        Pair<NamespaceBundles, List<NamespaceBundle>> splitBundles = factory.splitBundles(bundle, numberSplitBundles);
        // existing_no_bundles(1) +
        // additional_new_split_bundle(4) -
        // parent_target_bundle(1)
        int totalExpectedSplitBundles = bundles.getBundles().size() + numberSplitBundles - 1;
        validateSplitBundlesRange(bundles.getFullBundle(), splitBundles.getRight());
        assertEquals(totalExpectedSplitBundles, splitBundles.getLeft().getBundles().size());

        // (2) split in 4: first bundle from above split bundles
        NamespaceBundleFactory utilityFactory = getNamespaceBundleFactory();
        NamespaceBundles bundles2 = splitBundles.getLeft();
        NamespaceBundle testChildBundle = bundles2.getBundles().get(0);

        Pair<NamespaceBundles, List<NamespaceBundle>> splitChildBundles = splitBundlesUtilFactory(utilityFactory,
                nsname, bundles2, testChildBundle, numberSplitBundles);
        // existing_no_bundles(4) +
        // additional_new_split_bundle(4) -
        // parent_target_bundle(1)
        totalExpectedSplitBundles = bundles2.getBundles().size() + numberSplitBundles - 1;
        validateSplitBundlesRange(testChildBundle, splitChildBundles.getRight());
        assertEquals(totalExpectedSplitBundles, splitChildBundles.getLeft().getBundles().size());

        // (3) split in 3: second bundle from above split bundles
        NamespaceBundle testChildBundl2 = bundles2.getBundles().get(1);

        Pair<NamespaceBundles, List<NamespaceBundle>> splitChildBundles2 = splitBundlesUtilFactory(utilityFactory,
                nsname, bundles2, testChildBundl2, 3);
        // existing_no_bundles(4) +
        // additional_new_split_bundle(3) -
        // parent_target_bundle(1)
        totalExpectedSplitBundles = bundles2.getBundles().size() + 3 - 1;
        validateSplitBundlesRange(testChildBundl2, splitChildBundles2.getRight());
        assertEquals(totalExpectedSplitBundles, splitChildBundles2.getLeft().getBundles().size());
    }

    @Test
    public void testSplitBundleInTwo() throws Exception {
        final int NO_BUNDLES = 2;
        NamespaceName nsname = NamespaceName.get("pulsar/global/ns1");
        TopicName topicName = TopicName.get("persistent://pulsar/global/ns1/topic-1");
        NamespaceBundles bundles = factory.getBundles(nsname);
        NamespaceBundle bundle = bundles.findBundle(topicName);
        // (1) split : [0x00000000,0xffffffff] => [0x00000000_0x7fffffff,0x7fffffff_0xffffffff]
        Pair<NamespaceBundles, List<NamespaceBundle>> splitBundles = factory.splitBundles(bundle, NO_BUNDLES);
        assertNotNull(splitBundles);
        assertBundleDivideInTwo(bundle, splitBundles.getRight(), NO_BUNDLES);

        // (2) split: [0x00000000,0x7fffffff] => [0x00000000_0x3fffffff,0x3fffffff_0x7fffffff],
        // [0x7fffffff,0xffffffff] => [0x7fffffff_0xbfffffff,0xbfffffff_0xffffffff]
        NamespaceBundleFactory utilityFactory = getNamespaceBundleFactory();
        assertBundles(utilityFactory, nsname, bundle, splitBundles, NO_BUNDLES);

        // (3) split: [0x00000000,0x3fffffff] => [0x00000000_0x1fffffff,0x1fffffff_0x3fffffff],
        // [0x3fffffff,0x7fffffff] => [0x3fffffff_0x5fffffff,0x5fffffff_0x7fffffff]
        Pair<NamespaceBundles, List<NamespaceBundle>> splitChildBundles = splitBundlesUtilFactory(utilityFactory,
                nsname, splitBundles.getLeft(), splitBundles.getRight().get(0), NO_BUNDLES);
        assertBundles(utilityFactory, nsname, splitBundles.getRight().get(0), splitChildBundles, NO_BUNDLES);

        // (4) split: [0x7fffffff,0xbfffffff] => [0x7fffffff_0x9fffffff,0x9fffffff_0xbfffffff],
        // [0xbfffffff,0xffffffff] => [0xbfffffff_0xdfffffff,0xdfffffff_0xffffffff]
        splitChildBundles = splitBundlesUtilFactory(utilityFactory, nsname, splitBundles.getLeft(),
                splitBundles.getRight().get(1), NO_BUNDLES);
        assertBundles(utilityFactory, nsname, splitBundles.getRight().get(1), splitChildBundles, NO_BUNDLES);

    }

    private void validateSplitBundlesRange(NamespaceBundle fullBundle, List<NamespaceBundle> splitBundles) {
        assertNotNull(fullBundle);
        assertNotNull(splitBundles);
        Range<Long> fullRange = fullBundle.getKeyRange();
        Range<Long> span = splitBundles.get(0).getKeyRange();
        for (NamespaceBundle bundle : splitBundles) {
            span = span.span(bundle.getKeyRange());
        }
        assertEquals(span, fullRange);
    }

    @SuppressWarnings("unchecked")
    private Pair<NamespaceBundles, List<NamespaceBundle>> splitBundlesUtilFactory(NamespaceBundleFactory utilityFactory,
            NamespaceName nsname, NamespaceBundles bundles, NamespaceBundle targetBundle, int numBundles)
            throws Exception {
        Field bCacheField = NamespaceBundleFactory.class.getDeclaredField("bundlesCache");
        bCacheField.setAccessible(true);
        ((AsyncLoadingCache<NamespaceName, NamespaceBundles>) bCacheField.get(utilityFactory)).put(nsname,
                CompletableFuture.completedFuture(bundles));
        return utilityFactory.splitBundles(targetBundle, numBundles);
    }

    private void assertBundles(NamespaceBundleFactory utilityFactory, NamespaceName nsname, NamespaceBundle bundle,
            Pair<NamespaceBundles, List<NamespaceBundle>> splitBundles, int numBundles) throws Exception {

        NamespaceBundle bundle1 = splitBundles.getRight().get(0);
        NamespaceBundle bundle2 = splitBundles.getRight().get(1);

        NamespaceBundles nspaceBundles = splitBundles.getLeft();
        Pair<NamespaceBundles, List<NamespaceBundle>> bundle1Split = splitBundlesUtilFactory(utilityFactory, nsname,
                nspaceBundles, bundle1, numBundles);
        assertBundleDivideInTwo(bundle1, bundle1Split.getRight(), numBundles);

        Pair<NamespaceBundles, List<NamespaceBundle>> bundle2Split = splitBundlesUtilFactory(utilityFactory, nsname,
                nspaceBundles, bundle2, numBundles);
        assertBundleDivideInTwo(bundle2, bundle2Split.getRight(), numBundles);

    }

    private void assertBundleDivideInTwo(NamespaceBundle bundle, List<NamespaceBundle> bundles, int numBundles) {
        assertEquals(bundles.size(), 2);
        String[] range = bundle.getBundleRange().split("_");
        long lower = Long.decode(range[0]);
        long upper = Long.decode(range[1]);
        long middle = ((upper - lower) / 2) + lower;

        String lRange = String.format("0x%08x_0x%08x", lower, middle);
        String uRange = String.format("0x%08x_0x%08x", middle, upper);
        assertEquals(lRange, bundles.get(0).getBundleRange());
        assertEquals(uRange, bundles.get(1).getBundleRange());
        log.info("[{},{}] => [{},{}]", range[0], range[1], lRange, uRange);
    }

    private static final Logger log = LoggerFactory.getLogger(NamespaceBundlesTest.class);
}
