/*
 * 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.cassandra.dht.tokenallocator;

import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Random;
import java.util.Set;

import org.apache.commons.math3.stat.descriptive.SummaryStatistics;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import org.apache.cassandra.CassandraTestBase;
import org.apache.cassandra.CassandraTestBase.DDDaemonInitialization;
import org.apache.cassandra.CassandraTestBase.UseMurmur3Partitioner;
import org.apache.cassandra.ServerTestUtils;
import org.apache.cassandra.dht.IPartitioner;
import org.apache.cassandra.dht.Token;
import org.apache.cassandra.distributed.test.log.ClusterMetadataTestHelper;
import org.apache.cassandra.exceptions.ConfigurationException;
import org.apache.cassandra.locator.AbstractReplicationStrategy;
import org.apache.cassandra.locator.InetAddressAndPort;
import org.apache.cassandra.locator.SimpleLocationProvider;
import org.apache.cassandra.schema.KeyspaceMetadata;
import org.apache.cassandra.schema.KeyspaceParams;
import org.apache.cassandra.service.StorageService;
import org.apache.cassandra.tcm.ClusterMetadata;
import org.apache.cassandra.tcm.ClusterMetadataService;
import org.apache.cassandra.tcm.StubClusterMetadataService;
import org.apache.cassandra.tcm.membership.NodeId;
import org.apache.cassandra.utils.FBUtilities;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.fail;

@DDDaemonInitialization
@UseMurmur3Partitioner
public class TokenAllocationTest extends CassandraTestBase
{
    static Random rand = new Random(1);

    @Before
    public void before() throws ConfigurationException
    {
        ClusterMetadataService.setInstance(ClusterMetadataTestHelper.syncInstanceForTest());
        ClusterMetadataService.instance().log().unsafeBootstrapForTesting(FBUtilities.getBroadcastAddressAndPort());
    }

    @Before
    public void after() throws ConfigurationException
    {
        ClusterMetadataService.unsetInstance();
    }

    private static TokenAllocation createForTest(ClusterMetadata metadata, int replicas, int numTokens)
    {
        return TokenAllocation.create(metadata.locator.local().datacenter, metadata, replicas, numTokens);
    }

    @Test
    public void testAllocateTokensForKeyspace() throws UnknownHostException
    {
        int vn = 16;
        String ks = "TokenAllocationTestKeyspace3";
        ClusterMetadataTestHelper.addOrUpdateKeyspace(KeyspaceMetadata.create(ks, KeyspaceParams.simple(5)));
        ClusterMetadata metadata = generateFakeEndpoints(10, vn);
        InetAddressAndPort addr = FBUtilities.getBroadcastAddressAndPort();
        allocateTokensForKeyspace(vn, ks, metadata, addr);
    }

    @Test
    public void testAllocateTokensForLocalRF() throws UnknownHostException
    {
        int vn = 16;
        int allocateTokensForLocalRf = 3;
        ClusterMetadata metadata = generateFakeEndpoints(10, vn);
        InetAddressAndPort addr = FBUtilities.getBroadcastAddressAndPort();
        allocateTokensForLocalReplicationFactor(vn, allocateTokensForLocalRf, metadata, addr);
    }

    private Collection<Token> allocateTokensForKeyspace(int vnodes, String keyspace, ClusterMetadata metadata, InetAddressAndPort addr)
    {
        AbstractReplicationStrategy rs = metadata.schema.getKeyspaces().get(keyspace).get().replicationStrategy;
        TokenAllocation tokenAllocation = TokenAllocation.create(metadata, rs, vnodes);
        return allocateAndVerify(vnodes, addr, tokenAllocation);
    }

    private void allocateTokensForLocalReplicationFactor(int vnodes, int rf, ClusterMetadata metadata, InetAddressAndPort addr)
    {
        TokenAllocation tokenAllocation = createForTest(metadata, rf, vnodes);
        allocateAndVerify(vnodes, addr, tokenAllocation);
    }

    private Collection<Token> allocateAndVerify(int vnodes, InetAddressAndPort addr, TokenAllocation tokenAllocation)
    {
        SummaryStatistics os = tokenAllocation.getAllocationRingOwnership(addr);
        Collection<Token> tokens = tokenAllocation.allocate(addr);
        assertEquals(vnodes, tokens.size());
        SummaryStatistics ns = tokenAllocation.getAllocationRingOwnership(addr);
        if (ns.getStandardDeviation() > os.getStandardDeviation())
        {
            fail(String.format("Token allocation unexpectedly increased standard deviation.\nStats before:\n%s\nStats after:\n%s", os, ns));
        }
        return tokens;
    }

    @Test
    public void testAllocateTokensNetworkStrategyOneRack() throws UnknownHostException
    {
        testAllocateTokensNetworkStrategy(1, 3);
    }

    @Test(expected = ConfigurationException.class)
    public void testAllocateTokensNetworkStrategyTwoRacks() throws UnknownHostException
    {
        testAllocateTokensNetworkStrategy(2, 3);
    }

    @Test
    public void testAllocateTokensNetworkStrategyThreeRacks() throws UnknownHostException
    {
        testAllocateTokensNetworkStrategy(3, 3);
    }

    @Test
    public void testAllocateTokensNetworkStrategyFiveRacks() throws UnknownHostException
    {
        testAllocateTokensNetworkStrategy(5, 3);
    }

    @Test
    public void testAllocateTokensNetworkStrategyOneRackOneReplica() throws UnknownHostException
    {
        testAllocateTokensNetworkStrategy(1, 1);
    }


    public void testAllocateTokensNetworkStrategy(int rackCount, int replicas) throws UnknownHostException
    {
        int vn = 16;
        String ks = "TokenAllocationTestNTSKeyspace" + rackCount + replicas;
        String dc = "1";
        String otherDc = "15";
        KeyspaceMetadata keyspace = KeyspaceMetadata.create(ks, KeyspaceParams.nts(dc, replicas, otherDc, "15"));

        // register these 2 nodes so that the DCs exist, otherwise the CREATE KEYSPACE will be rejected
        // but don't join them, we don't assign any tokens to these nodes
        ClusterMetadataTestHelper.register(InetAddressAndPort.getByName("127.1.0.99"), dc, Integer.toString(0));
        ClusterMetadataTestHelper.register(InetAddressAndPort.getByName("127.15.0.99"), otherDc, Integer.toString(0));
        ClusterMetadataTestHelper.addOrUpdateKeyspace(keyspace);

        for (int i = 0; i < rackCount; ++i)
            generateFakeEndpoints(10, vn, dc, Integer.toString(i));
        InetAddressAndPort addr = InetAddressAndPort.getByName("127." + dc + ".0.99");
        allocateTokensForKeyspace(vn, ks, ClusterMetadata.current(), addr);
        // Note: Not matching replication factor in second datacentre, but this should not affect us.
    }

    @Test
    public void testAllocateTokensRfEqRacks() throws UnknownHostException
    {
        int vn = 8;
        int replicas = 3;
        int rackCount = replicas;
        String ks = "TokenAllocationTestNTSKeyspaceRfEqRack";
        String dc = "1";
        String otherDc = "15";
        KeyspaceMetadata keyspace = KeyspaceMetadata.create(ks, KeyspaceParams.nts(dc, replicas, otherDc, "15"));

        // register these 2 nodes so that the DCs exist, otherwise the CREATE KEYSPACE will be rejected
        // but don't join them, we don't assign any tokens to these nodes
        ClusterMetadataTestHelper.register(InetAddressAndPort.getByName("127.1.0.255"), dc, Integer.toString(0));
        ClusterMetadataTestHelper.register(InetAddressAndPort.getByName("127.15.0.255"), otherDc, Integer.toString(0));
        ClusterMetadataTestHelper.addOrUpdateKeyspace(keyspace);

        int base = 5;
        for (int i = 0; i < rackCount; ++i)
            generateFakeEndpoints(base << i, vn, dc, Integer.toString(i));     // unbalanced racks

        int cnt = 5;
        for (int i = 0; i < cnt; ++i)
        {
            InetAddressAndPort endpoint = InetAddressAndPort.getByName("127." + dc + ".0." + (99 + i));
            ClusterMetadataTestHelper.register(endpoint, dc, Integer.toString(0));
            Collection<Token> tokens = allocateTokensForKeyspace(vn, ks, ClusterMetadata.current(), endpoint);
            ClusterMetadataTestHelper.join(endpoint, new HashSet<>(tokens));
        }

        double target = 1.0 / (base + cnt);
        double permittedOver = 1.0 / (2 * vn + 1) + 0.01;

        Map<InetAddress, Float> ownership = StorageService.instance.effectiveOwnership(ks);
        boolean failed = false;
        for (Map.Entry<InetAddress, Float> o : ownership.entrySet())
        {
            int rack = o.getKey().getAddress()[2];
            if (rack != 0)
                continue;

            System.out.format("Node %s owns %f ratio to optimal %.2f\n", o.getKey(), o.getValue(), o.getValue() / target);
            if (o.getValue() / target > 1 + permittedOver)
                failed = true;
        }
        Assert.assertFalse(String.format("One of the nodes in the rack has over %.2f%% overutilization.", permittedOver * 100), failed);
    }

    /**
     * TODO: This scenario isn't supported very well. Investigate a multi-keyspace version of the algorithm.
     */
    @Test
    public void testAllocateTokensMultipleKeyspaces() throws UnknownHostException
    {
        final int TOKENS = 16;

        ClusterMetadataService.unsetInstance();
        ClusterMetadataService.setInstance(StubClusterMetadataService.forClientTools());
        generateFakeEndpoints(10, TOKENS);
        // pre-register the nodes which we are going to allocate tokens to this is necessary as
        // Location info is obtained by the allocator from cluster metadata, not the snitch
        Map<InetAddressAndPort, NodeId> endpointToIdMap = new HashMap<>(10);
        for (int i=11; i<=20; ++i)
        {
            InetAddressAndPort endpoint = InetAddressAndPort.getByName("127.0.0." + (i + 1));
            NodeId id = ClusterMetadataTestHelper.register(endpoint, SimpleLocationProvider.LOCATION);
            endpointToIdMap.put(endpoint, id);
        }
        ClusterMetadata metadata = ClusterMetadata.current();

        TokenAllocation rf2Allocator = createForTest(metadata, 2, TOKENS);
        TokenAllocation rf3Allocator = createForTest(metadata, 3, TOKENS);

        SummaryStatistics rf2StatsBefore = rf2Allocator.getAllocationRingOwnership(FBUtilities.getBroadcastAddressAndPort());
        SummaryStatistics rf3StatsBefore = rf3Allocator.getAllocationRingOwnership(FBUtilities.getBroadcastAddressAndPort());

        TokenAllocation current = rf3Allocator;
        TokenAllocation next = rf2Allocator;

        for (Map.Entry<InetAddressAndPort, NodeId> entry : endpointToIdMap.entrySet())
        {
            InetAddressAndPort endpoint = entry.getKey();
            NodeId id = entry.getValue();
            Set<Token> tokens = new HashSet<>(current.allocate(endpoint));
            // Update tokens on next to verify ownership calculation below
            next.updateTokensForNode(id, tokens);
            TokenAllocation tmp = current;
            current = next;
            next = tmp;
        }

        SummaryStatistics rf2StatsAfter = rf2Allocator.getAllocationRingOwnership(FBUtilities.getBroadcastAddressAndPort());
        SummaryStatistics rf3StatsAfter = rf3Allocator.getAllocationRingOwnership(FBUtilities.getBroadcastAddressAndPort());

        verifyImprovement(rf2StatsBefore, rf2StatsAfter);
        verifyImprovement(rf3StatsBefore, rf3StatsAfter);
    }

    private void verifyImprovement(SummaryStatistics os, SummaryStatistics ns)
    {
        if (ns.getStandardDeviation() > os.getStandardDeviation())
        {
            fail(String.format("Token allocation unexpectedly increased standard deviation.\nStats before:\n%s\nStats after:\n%s", os, ns));
        }
    }

    private ClusterMetadata generateFakeEndpoints(int nodes, int vnodes) throws UnknownHostException
    {
        // This is used by tests which previously relied on the ServerTestUtils snitch impl which has
        // 127.0.0.4 in datacenter2 and every other node in datacenter1
        IPartitioner p = ClusterMetadata.current().tokenMap.partitioner();
        for (int i = 1; i <= nodes; i++)
        {
            // leave .1 for myEndpoint
            InetAddressAndPort addr = InetAddressAndPort.getByName("127.0.0." + (i + 1));
            String dc = (i + 1 == 4) ? ServerTestUtils.DATA_CENTER_REMOTE : ServerTestUtils.DATA_CENTER;
            ClusterMetadataTestHelper.register(addr, dc, ServerTestUtils.RACK1);
            Set<Token> tokens = new HashSet<>(vnodes);
            for (int j = 0; j < vnodes; ++j)
                tokens.add(p.getRandomToken(rand));
            ClusterMetadataTestHelper.join(addr, tokens);
        }
        // Register (but don't join/allocate tokens for) 127.0.0.1)
        ClusterMetadataTestHelper.register(FBUtilities.getBroadcastAddressAndPort(),
                                           ServerTestUtils.DATA_CENTER,
                                           ServerTestUtils.RACK1);
        return ClusterMetadata.current();
    }

    private ClusterMetadata generateFakeEndpoints(int nodes, int vnodes, String dc, String rack) throws UnknownHostException
    {
        System.out.printf("Adding %d nodes to dc=%s, rack=%s.%n", nodes, dc, rack);
        IPartitioner p = ClusterMetadata.current().tokenMap.partitioner();

        for (int i = 1; i <= nodes; i++)
        {
            // leave .1 for myEndpoint
            InetAddressAndPort addr = InetAddressAndPort.getByName("127." + dc + '.' + rack + '.' + (i + 1));
            ClusterMetadataTestHelper.register(addr, dc, rack);
            Set<Token> tokens = new HashSet<>(vnodes);
            for (int j = 0; j < vnodes; ++j)
                tokens.add(p.getRandomToken(rand));
            ClusterMetadataTestHelper.join(addr, tokens);
        }
        return ClusterMetadata.current();
    }
}
