/*
 * 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.geode.internal.cache.control;

import static org.apache.geode.cache.PartitionAttributesFactory.GLOBAL_MAX_BUCKETS_DEFAULT;
import static org.apache.geode.management.runtime.RegionRedundancyStatus.RedundancyStatus.NOT_SATISFIED;
import static org.apache.geode.management.runtime.RegionRedundancyStatus.RedundancyStatus.NO_REDUNDANT_COPIES;
import static org.apache.geode.management.runtime.RegionRedundancyStatus.RedundancyStatus.SATISFIED;
import static org.apache.geode.management.runtime.RestoreRedundancyResults.Status.FAILURE;
import static org.apache.geode.management.runtime.RestoreRedundancyResults.Status.SUCCESS;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.Matchers.greaterThan;
import static org.hamcrest.Matchers.lessThanOrEqualTo;
import static org.hamcrest.Matchers.not;
import static org.hamcrest.Matchers.nullValue;
import static org.junit.Assert.assertThat;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.stream.IntStream;

import junitparams.JUnitParamsRunner;
import junitparams.Parameters;
import junitparams.naming.TestCaseName;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;

import org.apache.geode.cache.Cache;
import org.apache.geode.cache.Region;
import org.apache.geode.cache.RegionShortcut;
import org.apache.geode.cache.control.ResourceManager;
import org.apache.geode.internal.cache.PartitionAttributesImpl;
import org.apache.geode.internal.cache.PartitionedRegion;
import org.apache.geode.management.runtime.RegionRedundancyStatus;
import org.apache.geode.management.runtime.RestoreRedundancyResults;
import org.apache.geode.test.dunit.SerializableRunnableIF;
import org.apache.geode.test.dunit.rules.ClusterStartupRule;
import org.apache.geode.test.dunit.rules.MemberVM;

@RunWith(JUnitParamsRunner.class)
public class RestoreRedundancyOperationDUnitTest {
  private List<MemberVM> servers;
  private static final int SERVERS_TO_START = 3;
  private static final String PARENT_REGION_NAME = "parentColocatedRegion";
  private static final String CHILD_REGION_NAME = "childColocatedRegion";
  private static final int DESIRED_REDUNDANCY_COPIES = 2;
  private static final String LOW_REDUNDANCY_REGION_NAME = "lowRedundancyRegion";
  private static final int LOW_REDUNDANCY_COPIES = 1;
  private static final int ENTRIES = 5 * GLOBAL_MAX_BUCKETS_DEFAULT;

  @Rule
  public ClusterStartupRule cluster = new ClusterStartupRule();

  @Before
  public void startUp() {
    MemberVM locator = cluster.startLocatorVM(0);
    int locatorPort = locator.getPort();
    servers = new ArrayList<>();
    IntStream.range(0, SERVERS_TO_START)
        .forEach(i -> servers.add(cluster.startServerVM(i + 1, locatorPort)));

    // Create the regions on server1 and populate with data
    servers.get(0).invoke(() -> {
      Collection<Region<Object, Object>> regions = createRegions();
      regions.forEach(
          region -> IntStream.range(0, ENTRIES).forEach(i -> region.put("key" + i, "value" + i)));
    });

    // Create regions on other servers but do not populate with data
    servers.stream().skip(1).forEach(
        s -> s.invoke((SerializableRunnableIF) RestoreRedundancyOperationDUnitTest::createRegions));

    // Confirm that redundancy is impaired and primaries unbalanced for all regions on all members
    servers.forEach(s -> s.invoke(() -> {
      assertRedundancyStatus(PARENT_REGION_NAME, false);
      assertRedundancyStatus(CHILD_REGION_NAME, false);
      assertRedundancyStatus(LOW_REDUNDANCY_REGION_NAME, false);
      assertPrimariesBalanced(PARENT_REGION_NAME, SERVERS_TO_START, false);
      assertPrimariesBalanced(CHILD_REGION_NAME, SERVERS_TO_START, false);
      assertPrimariesBalanced(LOW_REDUNDANCY_REGION_NAME, SERVERS_TO_START, false);
    }));
  }

  @Test
  public void statsAreUpdatedWhenRestoreRedundancyIsCalled() {
    servers.get(0).invoke(() -> {
      restoreRedundancyAndGetResults(null, null, true);

      ResourceManagerStats stats = Objects.requireNonNull(ClusterStartupRule.getCache())
          .getInternalResourceManager().getStats();

      assertThat(stats.getRestoreRedundanciesInProgress(), equalTo(0L));
      assertThat(stats.getRestoreRedundanciesCompleted(), equalTo(1L));
      assertThat(stats.getRestoreRedundancyTime(), greaterThan(0L));
    });
  }

  @Test
  public void redundancyIsRecoveredAndPrimariesBalancedWhenRestoreRedundancyIsCalledWithNoIncludedOrExcludedRegions() {
    servers.get(0).invoke(() -> {
      RestoreRedundancyResults results = restoreRedundancyAndGetResults(null, null, true);
      assertThat(results.getRegionOperationStatus(), is(SUCCESS));
      assertThat(results.getTotalPrimaryTransfersCompleted() > 0, is(true));
      assertThat(results.getTotalPrimaryTransferTime() > 0, is(true));
      assertThat(results.getRegionResult(PARENT_REGION_NAME).getStatus(), is(SATISFIED));
      assertThat(results.getRegionResult(CHILD_REGION_NAME).getStatus(), is(SATISFIED));
      assertThat(results.getRegionResult(LOW_REDUNDANCY_REGION_NAME).getStatus(), is(SATISFIED));
    });

    servers.forEach(s -> s.invoke(() -> {
      assertRedundancyStatus(PARENT_REGION_NAME, true);
      assertRedundancyStatus(CHILD_REGION_NAME, true);
      assertRedundancyStatus(LOW_REDUNDANCY_REGION_NAME, true);
      assertPrimariesBalanced(PARENT_REGION_NAME, SERVERS_TO_START, true);
      assertPrimariesBalanced(CHILD_REGION_NAME, SERVERS_TO_START, true);
      assertPrimariesBalanced(LOW_REDUNDANCY_REGION_NAME, SERVERS_TO_START, true);
    }));
  }

  @Test
  public void redundancyIsRecoveredAndPrimariesNotBalancedWhenRestoreRedundancyIsCalledWithReassignPrimariesFalse() {
    servers.get(0).invoke(() -> {
      RestoreRedundancyResults results = restoreRedundancyAndGetResults(null, null, false);

      assertThat(results.getRegionOperationStatus(), is(SUCCESS));
      assertThat(results.getTotalPrimaryTransfersCompleted(), is(0));
      assertThat(results.getTotalPrimaryTransferTime(), is(0L));
      assertThat(results.getRegionResult(PARENT_REGION_NAME).getStatus(), is(SATISFIED));
      assertThat(results.getRegionResult(CHILD_REGION_NAME).getStatus(), is(SATISFIED));
      assertThat(results.getRegionResult(LOW_REDUNDANCY_REGION_NAME).getStatus(), is(SATISFIED));
    });

    servers.forEach(s -> s.invoke(() -> {
      assertRedundancyStatus(PARENT_REGION_NAME, true);
      assertRedundancyStatus(CHILD_REGION_NAME, true);
      assertRedundancyStatus(LOW_REDUNDANCY_REGION_NAME, true);
      assertPrimariesBalanced(PARENT_REGION_NAME, SERVERS_TO_START, false);
      assertPrimariesBalanced(CHILD_REGION_NAME, SERVERS_TO_START, false);
      assertPrimariesBalanced(LOW_REDUNDANCY_REGION_NAME, SERVERS_TO_START, false);
    }));
  }

  @Test
  public void redundancyIsNotRecoveredAndPrimariesNotBalancedForExcludedNonColocatedRegion() {
    servers.get(0).invoke(() -> {
      RestoreRedundancyResults results = restoreRedundancyAndGetResults(null,
          Collections.singleton(LOW_REDUNDANCY_REGION_NAME), true);

      assertThat(results.getRegionOperationStatus(), is(SUCCESS));
      assertThat(results.getRegionResult(PARENT_REGION_NAME).getStatus(), is(SATISFIED));
      assertThat(results.getRegionResult(CHILD_REGION_NAME).getStatus(), is(SATISFIED));
      assertThat(results.getRegionResult(LOW_REDUNDANCY_REGION_NAME), nullValue());
    });

    servers.forEach(s -> s.invoke(() -> {
      assertRedundancyStatus(PARENT_REGION_NAME, true);
      assertRedundancyStatus(CHILD_REGION_NAME, true);
      assertRedundancyStatus(LOW_REDUNDANCY_REGION_NAME, false);
      assertPrimariesBalanced(PARENT_REGION_NAME, SERVERS_TO_START, true);
      assertPrimariesBalanced(CHILD_REGION_NAME, SERVERS_TO_START, true);
      assertPrimariesBalanced(LOW_REDUNDANCY_REGION_NAME, SERVERS_TO_START, false);
    }));
  }

  @Test
  @Parameters(method = "getIncludeAndExclude")
  @TestCaseName("[{index}] {method}: Include={0}, Exclude={1}")
  public void redundancyIsRecoveredAndPrimariesBalancedForAllColocatedRegionsWhenAtLeastOneIsIncluded(
      String includeRegion, String excludeRegion) {
    servers.get(0).invoke(() -> {
      Set<String> includeSet = includeRegion == null ? null : Collections.singleton(includeRegion);
      Set<String> excludeSet = excludeRegion == null ? null : Collections.singleton(excludeRegion);

      RestoreRedundancyResults results =
          restoreRedundancyAndGetResults(includeSet, excludeSet, true);

      assertThat(results.getRegionOperationStatus(), is(SUCCESS));
      assertThat(results.getRegionResult(PARENT_REGION_NAME).getStatus(), is(SATISFIED));
      assertThat(results.getRegionResult(CHILD_REGION_NAME).getStatus(), is(SATISFIED));
    });

    servers.forEach(s -> s.invoke(() -> {
      assertRedundancyStatus(PARENT_REGION_NAME, true);
      assertRedundancyStatus(CHILD_REGION_NAME, true);
      assertPrimariesBalanced(PARENT_REGION_NAME, SERVERS_TO_START, true);
      assertPrimariesBalanced(CHILD_REGION_NAME, SERVERS_TO_START, true);
    }));
  }

  @Test
  public void restoringRedundancyWithoutEnoughServersToFullySatisfyRedundancyShouldReturnFailureStatusAndBalancePrimaries() {
    servers.remove(servers.size() - 1).stop();
    int activeServers = servers.size();

    servers.get(0).invoke(() -> {
      RestoreRedundancyResults results = restoreRedundancyAndGetResults(null, null, true);
      assertThat(results.getRegionOperationStatus(), is(FAILURE));
      assertThat(results.getRegionResult(PARENT_REGION_NAME).getStatus(), is(NOT_SATISFIED));
      assertThat(results.getRegionResult(CHILD_REGION_NAME).getStatus(), is(NOT_SATISFIED));
      assertThat(results.getRegionResult(LOW_REDUNDANCY_REGION_NAME).getStatus(), is(SATISFIED));
    });

    servers.forEach(s -> s.invoke(() -> {
      assertRedundancyStatus(PARENT_REGION_NAME, false);
      assertRedundancyStatus(CHILD_REGION_NAME, false);
      assertRedundancyStatus(LOW_REDUNDANCY_REGION_NAME, true);
      assertPrimariesBalanced(PARENT_REGION_NAME, activeServers, true);
      assertPrimariesBalanced(CHILD_REGION_NAME, activeServers, true);
      assertPrimariesBalanced(LOW_REDUNDANCY_REGION_NAME, activeServers, true);
    }));
  }

  @Test
  public void restoringRedundancyWithoutEnoughServersToCreateAnyRedundantCopyShouldReturnFailureStatus() {
    // Stop the last two servers in the list and remove them from the list, leaving us with one
    // server
    servers.remove(servers.size() - 1).stop();
    servers.remove(servers.size() - 1).stop();

    assertThat(servers.size(), is(1));

    servers.get(0).invoke(() -> {
      RestoreRedundancyResults results = restoreRedundancyAndGetResults(null, null, true);
      assertThat(results.getRegionOperationStatus(), is(FAILURE));
      assertThat(results.getRegionResult(PARENT_REGION_NAME).getStatus(), is(NO_REDUNDANT_COPIES));
      assertThat(results.getRegionResult(CHILD_REGION_NAME).getStatus(), is(NO_REDUNDANT_COPIES));
      assertThat(results.getRegionResult(LOW_REDUNDANCY_REGION_NAME).getStatus(),
          is(NO_REDUNDANT_COPIES));
      assertRedundancyStatus(PARENT_REGION_NAME, false);
      assertRedundancyStatus(CHILD_REGION_NAME, false);
      assertRedundancyStatus(LOW_REDUNDANCY_REGION_NAME, false);
    });
  }

  @Test
  public void redundancyStatusReportsCorrectRedundancyInformation() {
    getRedundancyStatusAndAssertIfSatisfied(FAILURE, NO_REDUNDANT_COPIES);

    // Restore redundancy to confirm that the result of the redundancyStatus call has changed to
    // reflect the new state of the system
    servers.get(0).invoke(() -> {
      restoreRedundancyAndGetResults(null, null, true);
    });

    getRedundancyStatusAndAssertIfSatisfied(SUCCESS, SATISFIED);
  }

  private static RestoreRedundancyResults restoreRedundancyAndGetResults(
      Set<String> includeRegions, Set<String> excludeRegions, boolean shouldReassign)
      throws InterruptedException, ExecutionException {
    ResourceManager resourceManager =
        Objects.requireNonNull(ClusterStartupRule.getCache()).getResourceManager();
    CompletableFuture<RestoreRedundancyResults> redundancyOpFuture = resourceManager
        .createRestoreRedundancyOperation()
        .includeRegions(includeRegions)
        .excludeRegions(excludeRegions)
        .shouldReassignPrimaries(shouldReassign)
        .start();
    assertThat(resourceManager.getRestoreRedundancyFutures().size(), is(1));
    assertThat(resourceManager.getRestoreRedundancyFutures().contains(redundancyOpFuture),
        is(true));
    return redundancyOpFuture.get();
  }

  private static Collection<Region<Object, Object>> createRegions() {
    Collection<Region<Object, Object>> regions = new HashSet<>();
    PartitionAttributesImpl attributes = getAttributesWithRedundancy(DESIRED_REDUNDANCY_COPIES);
    regions.add(Objects.requireNonNull(ClusterStartupRule.getCache())
        .createRegionFactory(RegionShortcut.PARTITION).setPartitionAttributes(attributes)
        .create(PARENT_REGION_NAME));

    attributes.setColocatedWith(PARENT_REGION_NAME);
    regions.add(Objects.requireNonNull(ClusterStartupRule.getCache())
        .createRegionFactory(RegionShortcut.PARTITION).setPartitionAttributes(attributes)
        .create(CHILD_REGION_NAME));

    PartitionAttributesImpl lowRedundancyAttributes =
        getAttributesWithRedundancy(LOW_REDUNDANCY_COPIES);
    regions.add(Objects.requireNonNull(ClusterStartupRule.getCache())
        .createRegionFactory(RegionShortcut.PARTITION)
        .setPartitionAttributes(lowRedundancyAttributes).create(LOW_REDUNDANCY_REGION_NAME));

    return regions;
  }

  private static PartitionAttributesImpl getAttributesWithRedundancy(int desiredRedundancy) {
    PartitionAttributesImpl attributes = new PartitionAttributesImpl();
    attributes.setRedundantCopies(desiredRedundancy);
    attributes.setRecoveryDelay(-1);
    attributes.setStartupRecoveryDelay(-1);
    attributes.setTotalNumBuckets(GLOBAL_MAX_BUCKETS_DEFAULT);
    return attributes;
  }

  private static void assertRedundancyStatus(String regionName, boolean shouldBeSatisfied) {
    Cache cache = Objects.requireNonNull(ClusterStartupRule.getCache());

    PartitionedRegion region = (PartitionedRegion) cache.getRegion(regionName);

    String message =
        "Expecting redundancy to " + (shouldBeSatisfied ? "" : "not") + " be satisfied";
    assertThat(message, region.getRedundancyProvider().isRedundancyImpaired(),
        is(!shouldBeSatisfied));
  }

  private static void assertPrimariesBalanced(String regionName, int numberOfServers,
      boolean shouldBeBalanced) {
    Cache cache = Objects.requireNonNull(ClusterStartupRule.getCache());

    PartitionedRegion region = (PartitionedRegion) cache.getRegion(regionName);
    int primariesOnServer = region.getLocalPrimaryBucketsListTestOnly().size();
    // Add one to account for integer rounding errors when dividing
    int expectedPrimaries = 1 + GLOBAL_MAX_BUCKETS_DEFAULT / numberOfServers;
    // Because of the way reassigning primaries works, it is sometimes only possible to get the
    // difference between the most loaded member and the least loaded member to be 2, not 1 as would
    // be the case for perfect balance
    String message = "Primaries should be balanced, but expectedPrimaries:actualPrimaries = "
        + expectedPrimaries + ":" + primariesOnServer;
    if (shouldBeBalanced) {
      assertThat(message, Math.abs(primariesOnServer - expectedPrimaries),
          is(lessThanOrEqualTo(2)));
    } else {
      assertThat("Primaries should not be balanced",
          Math.abs(primariesOnServer - expectedPrimaries), is(not(lessThanOrEqualTo(2))));
    }
  }

  private void getRedundancyStatusAndAssertIfSatisfied(
      RestoreRedundancyResults.Status expectedResultStatus,
      RegionRedundancyStatus.RedundancyStatus expectedRedundancyStatus) {
    servers.forEach(s -> s.invoke(() -> {
      RestoreRedundancyResults results = Objects.requireNonNull(ClusterStartupRule.getCache())
          .getResourceManager()
          .createRestoreRedundancyOperation()
          .redundancyStatus();
      assertThat(results.getRegionOperationStatus(), is(expectedResultStatus));
      assertThat(results.getRegionResult(PARENT_REGION_NAME).getStatus(), is(
          expectedRedundancyStatus));
      assertThat(results.getRegionResult(CHILD_REGION_NAME).getStatus(), is(
          expectedRedundancyStatus));
      assertThat(results.getRegionResult(LOW_REDUNDANCY_REGION_NAME).getStatus(),
          is(expectedRedundancyStatus));
    }));
  }

  @SuppressWarnings("unused")
  private Object[] getIncludeAndExclude() {
    return new Object[] {
        new Object[] {PARENT_REGION_NAME, null},
        new Object[] {CHILD_REGION_NAME, null},
        new Object[] {CHILD_REGION_NAME, PARENT_REGION_NAME},
        new Object[] {PARENT_REGION_NAME, CHILD_REGION_NAME},
        new Object[] {null, PARENT_REGION_NAME},
        new Object[] {null, CHILD_REGION_NAME}
    };
  }
}