/*
 * 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.stream.storage.impl.sc;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertTrue;

import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.LongStream;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.net.BookieId;
import org.apache.bookkeeper.net.BookieSocketAddress;
import org.apache.bookkeeper.stream.proto.cluster.ClusterAssignmentData;
import org.apache.bookkeeper.stream.proto.cluster.ClusterMetadata;
import org.apache.bookkeeper.stream.proto.cluster.ServerAssignmentData;
import org.apache.bookkeeper.stream.storage.impl.sc.DefaultStorageContainerController.ServerAssignmentDataComparator;
import org.apache.commons.lang3.tuple.Pair;
import org.junit.Test;

/**
 * Unit test {@link DefaultStorageContainerController}.
 */
@Slf4j
public class DefaultStorageContainerControllerTest {

    private static final int NUM_STORAGE_CONTAINERS = 32;

    private final ClusterMetadata clusterMetadata;
    private final StorageContainerController controller;
    private final ClusterAssignmentData currentAssignment;

    public DefaultStorageContainerControllerTest() {
        this.controller = new DefaultStorageContainerController();
        this.clusterMetadata = ClusterMetadata.newBuilder()
            .setNumStorageContainers(NUM_STORAGE_CONTAINERS)
            .build();
        this.currentAssignment = ClusterAssignmentData.newBuilder()
            .putServers("default-server", ServerAssignmentData.newBuilder()
                .addContainers(0L)
                .addContainers(1L)
                .addContainers(3L)
                .build())
            .build();
    }

    @Test
    public void testServerAssignmentDataComparator() {
        ServerAssignmentDataComparator comparator = new ServerAssignmentDataComparator();

        LinkedList<Long> serverList1 = new LinkedList<>();
        serverList1.add(1L);
        LinkedList<Long> serverList2 = new LinkedList<>();
        serverList2.add(2L);
        serverList2.add(3L);

        BookieId address1 = new BookieSocketAddress("127.0.0.1", 4181).toBookieId();
        BookieId address2 = new BookieSocketAddress("127.0.0.1", 4182).toBookieId();

        Pair<BookieId, LinkedList<Long>> pair1 = Pair.of(address1, serverList1);
        Pair<BookieId, LinkedList<Long>> pair2 = Pair.of(address1, serverList2);
        Pair<BookieId, LinkedList<Long>> pair3 = Pair.of(address2, serverList2);

        assertEquals(-1, comparator.compare(pair1, pair2));
        assertEquals(-1, comparator.compare(pair1, pair2));
        assertEquals(
            Integer.compare(address1.hashCode(), address2.hashCode()),
            comparator.compare(pair2, pair3));
    }

    @Test
    public void testComputeIdealStateEmptyCluster() {
        assertSame(
            currentAssignment,
            controller.computeIdealState(
                clusterMetadata,
                currentAssignment,
                Collections.emptySet()));
    }

    private static Set<BookieId> newCluster(int numServers) {
        Set<BookieId> cluster = IntStream.range(0, numServers)
            .mapToObj(idx -> new BookieSocketAddress("127.0.0.1", 4181 + idx).toBookieId())
            .collect(Collectors.toSet());
        return ImmutableSet.copyOf(cluster);
    }

    private static Set<BookieId> newCluster(int numServers, int startServerIdx) {
        Set<BookieId> cluster = IntStream.range(0, numServers)
            .mapToObj(idx -> new BookieSocketAddress("127.0.0.1", 4181 + startServerIdx + idx).toBookieId())
            .collect(Collectors.toSet());
        return ImmutableSet.copyOf(cluster);
    }

    private static void verifyAssignmentData(ClusterAssignmentData newAssignment,
                                             Set<BookieId> currentCluster,
                                             boolean isInitialIdealState)
            throws Exception {
        int numServers = currentCluster.size();

        assertEquals(numServers, newAssignment.getServersCount());
        Set<Long> assignedContainers = Sets.newHashSet();
        Set<BookieId> assignedServers = Sets.newHashSet();

        int numContainersPerServer = NUM_STORAGE_CONTAINERS / numServers;
        int serverIdx = 0;
        for (Map.Entry<String, ServerAssignmentData> entry : newAssignment.getServersMap().entrySet()) {
            log.info("Check assignment for server {} = {}", entry.getKey(), entry.getValue());

            BookieId address = BookieId.parse(entry.getKey());
            assignedServers.add(address);
            assertEquals(serverIdx + 1, assignedServers.size());

            ServerAssignmentData serverData = entry.getValue();
            assertEquals(numContainersPerServer, serverData.getContainersCount());
            List<Long> containers = Lists.newArrayList(serverData.getContainersList());
            Collections.sort(containers);
            assignedContainers.addAll(containers);

            if (isInitialIdealState) {
                long startContainerId = containers.get(0);
                for (int i = 0; i < containers.size(); i++) {
                    assertEquals(startContainerId + i * numServers, containers.get(i).longValue());
                }
            }
            ++serverIdx;
        }

        // each server should be assigned with equal number of containers
        assertTrue(Sets.difference(currentCluster, assignedServers).isEmpty());
        // all containers should be assigned
        Set<Long> expectedContainers = LongStream.range(0L, NUM_STORAGE_CONTAINERS)
            .boxed()
            .collect(Collectors.toSet());
        assertTrue(Sets.difference(expectedContainers, assignedContainers).isEmpty());
    }

    private static void verifyAssignmentDataWhenHasMoreServers(ClusterAssignmentData newAssignment,
                                                               Set<BookieId> currentCluster)
            throws Exception {
        int numServers = currentCluster.size();

        assertEquals(numServers, newAssignment.getServersCount());
        Set<Long> assignedContainers = Sets.newHashSet();
        Set<BookieId> assignedServers = Sets.newHashSet();

        int numEmptyServers = 0;
        int numAssignedServers = 0;
        int serverIdx = 0;
        for (Map.Entry<String, ServerAssignmentData> entry : newAssignment.getServersMap().entrySet()) {
            log.info("Check assignment for server {} = {}", entry.getKey(), entry.getValue());

            BookieId address = BookieId.parse(entry.getKey());
            assignedServers.add(address);
            assertEquals(serverIdx + 1, assignedServers.size());

            ServerAssignmentData serverData = entry.getValue();
            if (serverData.getContainersCount() > 0) {
                assertEquals(1, serverData.getContainersCount());
                ++numAssignedServers;
            } else {
                ++numEmptyServers;
            }
            List<Long> containers = Lists.newArrayList(serverData.getContainersList());
            Collections.sort(containers);
            assignedContainers.addAll(containers);

            ++serverIdx;
        }

        assertEquals(numServers / 2, numEmptyServers);
        assertEquals(numServers / 2, numAssignedServers);

        // each server should be assigned with equal number of containers
        assertTrue(Sets.difference(currentCluster, assignedServers).isEmpty());
        // all containers should be assigned
        Set<Long> expectedContainers = LongStream.range(0L, NUM_STORAGE_CONTAINERS)
            .boxed()
            .collect(Collectors.toSet());
        assertTrue(Sets.difference(expectedContainers, assignedContainers).isEmpty());
    }

    @Test
    public void testComputeIdealStateFromEmptyAssignment() throws Exception {
        ClusterAssignmentData emptyAssignment = ClusterAssignmentData.newBuilder().build();

        int numServers = 8;
        Set<BookieId> currentCluster = newCluster(numServers);

        ClusterAssignmentData newAssignment = controller.computeIdealState(
            clusterMetadata,
            emptyAssignment,
            currentCluster);

        verifyAssignmentData(newAssignment, currentCluster, true);
    }

    @Test
    public void testComputeIdealStateIfClusterUnchanged() throws Exception {
        ClusterAssignmentData emptyAssignment = ClusterAssignmentData.newBuilder().build();

        int numServers = 8;
        Set<BookieId> currentCluster = newCluster(numServers);
        ClusterAssignmentData newAssignment = controller.computeIdealState(
            clusterMetadata,
            emptyAssignment,
            currentCluster);
        verifyAssignmentData(newAssignment, currentCluster, true);

        ClusterAssignmentData newAssignment2 = controller.computeIdealState(
            clusterMetadata,
            newAssignment,
            currentCluster);

        // the state should not change if cluster is unchanged.
        assertSame(newAssignment, newAssignment2);
    }

    @Test
    public void testComputeIdealStateWhenHostsRemoved() throws Exception {
        ClusterAssignmentData emptyAssignment = ClusterAssignmentData.newBuilder().build();

        int numServers = 8;
        Set<BookieId> currentCluster = newCluster(numServers);

        ClusterAssignmentData assignmentData = controller.computeIdealState(
            clusterMetadata,
            emptyAssignment,
            currentCluster);
        verifyAssignmentData(assignmentData, currentCluster, true);

        int newNumServers = 4;
        Set<BookieId> newCluster = newCluster(newNumServers);

        ClusterAssignmentData newAssignmentData = controller.computeIdealState(
            clusterMetadata,
            assignmentData,
            newCluster);
        verifyAssignmentData(newAssignmentData, newCluster, false);
    }

    @Test
    public void testComputeIdealStateWhenHostsAdded() throws Exception {
        ClusterAssignmentData emptyAssignment = ClusterAssignmentData.newBuilder().build();

        int numServers = 4;
        Set<BookieId> currentCluster = newCluster(numServers);

        ClusterAssignmentData assignmentData = controller.computeIdealState(
            clusterMetadata,
            emptyAssignment,
            currentCluster);
        verifyAssignmentData(assignmentData, currentCluster, true);

        int newNumServers = 8;
        Set<BookieId> newCluster = newCluster(newNumServers);

        ClusterAssignmentData newAssignmentData = controller.computeIdealState(
            clusterMetadata,
            assignmentData,
            newCluster);
        verifyAssignmentData(newAssignmentData, newCluster, false);
    }

    @Test
    public void testComputeIdealStateWhenHostsRemovedAdded() throws Exception {
        ClusterAssignmentData emptyAssignment = ClusterAssignmentData.newBuilder().build();

        int numServers = 4;
        Set<BookieId> currentCluster = newCluster(numServers);

        ClusterAssignmentData assignmentData = controller.computeIdealState(
            clusterMetadata,
            emptyAssignment,
            currentCluster);
        verifyAssignmentData(assignmentData, currentCluster, true);

        Set<BookieId> serversToAdd = newCluster(6, numServers);
        Set<BookieId> serversToRemove = newCluster(2);

        Set<BookieId> newCluster = Sets.newHashSet(currentCluster);
        newCluster.addAll(serversToAdd);
        serversToRemove.forEach(newCluster::remove);

        ClusterAssignmentData newAssignmentData = controller.computeIdealState(
            clusterMetadata,
            assignmentData,
            newCluster);
        verifyAssignmentData(newAssignmentData, newCluster, false);
    }

    @Test
    public void testComputeIdealStateWhenHasMoreServers() throws Exception {
        ClusterAssignmentData emptyAssignment = ClusterAssignmentData.newBuilder().build();

        int numServers = 2 * NUM_STORAGE_CONTAINERS;
        Set<BookieId> currentCluster = newCluster(numServers);

        ClusterAssignmentData assignmentData = controller.computeIdealState(
            clusterMetadata,
            emptyAssignment,
            currentCluster);
        verifyAssignmentDataWhenHasMoreServers(assignmentData, currentCluster);
    }

    @Test
    public void testComputeIdealStateWhenScaleToMoreServers() throws Exception {
        ClusterAssignmentData emptyAssignment = ClusterAssignmentData.newBuilder().build();

        int numServers = 4;
        Set<BookieId> currentCluster = newCluster(numServers);

        ClusterAssignmentData assignmentData = controller.computeIdealState(
            clusterMetadata,
            emptyAssignment,
            currentCluster);
        verifyAssignmentData(assignmentData, currentCluster, true);

        numServers = 2 * NUM_STORAGE_CONTAINERS;
        Set<BookieId> newCluster = newCluster(numServers);
        ClusterAssignmentData newAssignment = controller.computeIdealState(
            clusterMetadata,
            assignmentData,
            newCluster);
        verifyAssignmentDataWhenHasMoreServers(newAssignment, newCluster);
    }

}
