package voldemort.routing;

import java.util.HashMap;
import java.util.List;

import junit.framework.TestCase;
import voldemort.TestUtils;
import voldemort.cluster.Node;

import com.google.common.collect.Lists;

public class ConsistentToZoneRoutingTest extends TestCase {

    public void testBackwardsCompatibility() {
        backwardsCompatibility(10, 2, 3, 2, 4);
    }

    public void backwardsCompatibility(int numNodes,
                                       int numPartitionsPerNode,
                                       int numReplicas,
                                       int numZones,
                                       int numKeys) {

        List<Integer> partitions = Lists.newArrayList();
        for(int i = 0; i < numNodes * numPartitionsPerNode; i++)
            partitions.add(i);

        List<Node> nodes = randomNodesInOneZone(numNodes, partitions);
        ConsistentRoutingStrategy beforeStrategy = new ConsistentRoutingStrategy(nodes, numReplicas);

        List<Node> newNodes = doubleToMultiZone(nodes, numPartitionsPerNode, numZones);
        ZoneRoutingStrategy afterStrategy = new ZoneRoutingStrategy(newNodes,
                                                                    getTestZoneReplicationFactor(numReplicas,
                                                                                                 numZones),
                                                                    numReplicas * numZones);

        for(int i = 0; i < numKeys; i++) {
            byte[] randomKey = TestUtils.randomBytes(10);
            List<Integer> beforePartitions = beforeStrategy.getPartitionList(randomKey);
            List<Integer> afterPartitions = afterStrategy.getPartitionList(randomKey);
            beforePartitions.retainAll(afterPartitions);
            assertTrue(beforePartitions.size() > 0);
        }
    }

    public HashMap<Integer, Integer> getTestZoneReplicationFactor(int replicationFactorPerZone,
                                                                  int numZones) {
        HashMap<Integer, Integer> returnHashMap = new HashMap<Integer, Integer>();
        for(int i = 0; i < numZones; i++) {
            returnHashMap.put(i, replicationFactorPerZone);
        }
        return returnHashMap;
    }

    public List<Node> doubleToMultiZone(List<Node> nodes, int numPartitionsPerNode, int numZones) {
        if(numPartitionsPerNode % numZones != 0) {
            return null;
        }

        List<Node> returnedNodes = Lists.newArrayList();
        for(Node node: nodes) {
            List<Integer> partitions = node.getPartitionIds();

            int numPartitionsPerZone = partitions.size() / numZones;
            for(int i = 0; i < numZones; i++) {
                returnedNodes.add(node((i * nodes.size()) + node.getId(),
                                       i,
                                       partitions.subList(i * numPartitionsPerZone,
                                                          numPartitionsPerZone * (i + 1))));
            }
        }
        return returnedNodes;
    }

    public List<Node> randomNodesInOneZone(int numNodes, List<Integer> partitions) {
        if(partitions.size() % numNodes != 0) {
            return null;
        }

        int partitionsPerNode = partitions.size() / numNodes;
        List<Node> returnedNodes = Lists.newArrayList();
        for(int i = 0; i < numNodes; i++) {
            returnedNodes.add(node(i, 0, partitions.subList(partitionsPerNode * i,
                                                            partitionsPerNode * (i + 1))));
        }
        return returnedNodes;
    }

    private Node node(int id, int zoneId, List<Integer> list) {
        return new Node(id, "localhost", 8080, 6666, 6667, zoneId, list);
    }
}