/*
 * 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 com.datastax.driver.core;

import static com.datastax.driver.core.CreateCCM.TestMode.PER_METHOD;
import static org.testng.Assert.assertTrue;
import static org.testng.Assert.fail;

import com.datastax.driver.core.exceptions.NoHostAvailableException;
import com.datastax.driver.core.policies.ConstantReconnectionPolicy;
import com.datastax.driver.core.policies.ExponentialReconnectionPolicy;
import com.datastax.driver.core.policies.ReconnectionPolicy;
import org.testng.annotations.Test;

@CreateCCM(PER_METHOD)
@CCMConfig(dirtiesContext = true, createKeyspace = false)
@SuppressWarnings("unused")
public class ReconnectionPolicyTest extends AbstractPoliciesTest {

  private Cluster.Builder exponential() {
    return Cluster.builder()
        .withReconnectionPolicy(new ExponentialReconnectionPolicy(2 * 1000, 5 * 60 * 1000));
  }

  private Cluster.Builder constant() {
    return Cluster.builder().withReconnectionPolicy(new ConstantReconnectionPolicy(10 * 1000));
  }

  /*
   * Test the ExponentialReconnectionPolicy.
   */
  @Test(groups = "long")
  @CCMConfig(clusterProvider = "exponential")
  public void exponentialReconnectionPolicyTest() throws Throwable {

    // Ensure that ExponentialReconnectionPolicy is what we should be testing
    if (!(cluster().getConfiguration().getPolicies().getReconnectionPolicy()
        instanceof ExponentialReconnectionPolicy)) {
      fail("Set policy does not match retrieved policy.");
    }

    // Test basic getters
    ExponentialReconnectionPolicy reconnectionPolicy =
        (ExponentialReconnectionPolicy)
            cluster().getConfiguration().getPolicies().getReconnectionPolicy();
    assertTrue(reconnectionPolicy.getBaseDelayMs() == 2 * 1000);
    assertTrue(reconnectionPolicy.getMaxDelayMs() == 5 * 60 * 1000);

    // Test erroneous instantiations
    try {
      new ExponentialReconnectionPolicy(-1, 1);
      fail();
    } catch (IllegalArgumentException e) {
      // ok
    }

    try {
      new ExponentialReconnectionPolicy(1, -1);
      fail();
    } catch (IllegalArgumentException e) {
      // ok
    }

    try {
      new ExponentialReconnectionPolicy(-1, -1);
      fail();
    } catch (IllegalArgumentException e) {
      // ok
    }

    try {
      new ExponentialReconnectionPolicy(2, 1);
      fail();
    } catch (IllegalArgumentException e) {
      // ok
    }

    // Test nextDelays()
    ReconnectionPolicy.ReconnectionSchedule schedule =
        new ExponentialReconnectionPolicy(2 * 1000, 5 * 60 * 1000).newSchedule();
    assertTrue(schedule.nextDelayMs() == 2000);
    assertTrue(schedule.nextDelayMs() == 4000);
    assertTrue(schedule.nextDelayMs() == 8000);
    assertTrue(schedule.nextDelayMs() == 16000);
    assertTrue(schedule.nextDelayMs() == 32000);
    for (int i = 0; i < 64; ++i) schedule.nextDelayMs();
    assertTrue(schedule.nextDelayMs() == reconnectionPolicy.getMaxDelayMs());

    // Run integration test
    long restartTime = 2 + 4 + 8 + 2; // 16: 3 full cycles + 2 seconds
    long retryTime = 30; // 4th cycle start time
    long breakTime = 62; // time until next reconnection attempt

    // TODO: Try to sort out variance
    // reconnectionPolicyTest(restartTime, retryTime, breakTime);
  }

  /*
   * Test the ConstantReconnectionPolicy.
   */
  @Test(groups = "long")
  @CCMConfig(clusterProvider = "constant")
  public void constantReconnectionPolicyTest() throws Throwable {

    // Ensure that ConstantReconnectionPolicy is what we should be testing
    if (!(cluster().getConfiguration().getPolicies().getReconnectionPolicy()
        instanceof ConstantReconnectionPolicy)) {
      fail("Set policy does not match retrieved policy.");
    }

    // Test basic getters
    ConstantReconnectionPolicy reconnectionPolicy =
        (ConstantReconnectionPolicy)
            cluster().getConfiguration().getPolicies().getReconnectionPolicy();
    assertTrue(reconnectionPolicy.getConstantDelayMs() == 10 * 1000);

    // Test erroneous instantiations
    try {
      new ConstantReconnectionPolicy(-1);
      fail();
    } catch (IllegalArgumentException e) {
      // ok
    }

    // Test nextDelays()
    ReconnectionPolicy.ReconnectionSchedule schedule =
        new ConstantReconnectionPolicy(10 * 1000).newSchedule();
    assertTrue(schedule.nextDelayMs() == 10000);
    assertTrue(schedule.nextDelayMs() == 10000);
    assertTrue(schedule.nextDelayMs() == 10000);
    assertTrue(schedule.nextDelayMs() == 10000);
    assertTrue(schedule.nextDelayMs() == 10000);

    // Run integration test
    long restartTime = 32; // matches the above test
    long retryTime = 40; // 2nd cycle start time
    long breakTime = 10; // time until next reconnection attempt

    // TODO: Try to sort out variance
    // reconnectionPolicyTest(restartTime, retryTime, breakTime);
  }

  public void reconnectionPolicyTest(long restartTime, long retryTime, long breakTime)
      throws Throwable {
    createSchema(1);
    init(12);
    query(12);

    // Ensure a basic test works
    assertQueried(TestUtils.IP_PREFIX + '1', 12);
    resetCoordinators();
    ccm().forceStop(1);

    // Start timing and ensure that the node is down
    long startTime = 0;
    try {
      startTime = System.nanoTime() / 1000000000;
      query(12);
      fail("Test race condition where node has not shut off quickly enough.");
    } catch (NoHostAvailableException e) {
      // ok
    }

    long thisTime;
    boolean restarted = false;
    while (true) {
      thisTime = System.nanoTime() / 1000000000;

      // Restart node at restartTime
      if (!restarted && thisTime - startTime > restartTime) {
        ccm().start(1);
        restarted = true;
      }

      // Continue testing queries each second
      try {
        query(12);
        assertQueried(TestUtils.IP_PREFIX + '1', 12);
        resetCoordinators();

        // Ensure the time when the query completes successfully is what was expected
        assertTrue(
            retryTime - 2 < thisTime - startTime && thisTime - startTime < retryTime + 2,
            String.format(
                "Waited %s seconds instead an expected %s seconds wait",
                thisTime - startTime, retryTime));
      } catch (NoHostAvailableException e) {
        Thread.sleep(1000);
        continue;
      }

      Thread.sleep(breakTime);

      // The the same query once more, just to be sure
      query(12);
      assertQueried(TestUtils.IP_PREFIX + '1', 12);
      resetCoordinators();

      // Ensure the reconnection times reset
      ccm().forceStop(1);

      // Start timing and ensure that the node is down
      startTime = 0;
      try {
        startTime = System.nanoTime() / 1000000000;
        query(12);
        fail("Test race condition where node has not shut off quickly enough.");
      } catch (NoHostAvailableException e) {
        // ok
      }

      restarted = false;
      while (true) {
        thisTime = System.nanoTime() / 1000000000;

        // Restart node at restartTime
        if (!restarted && thisTime - startTime > restartTime) {
          ccm().start(1);
          restarted = true;
        }

        // Continue testing queries each second
        try {
          query(12);
          assertQueried(TestUtils.IP_PREFIX + '1', 12);
          resetCoordinators();

          // Ensure the time when the query completes successfully is what was expected
          assertTrue(
              retryTime - 2 < thisTime - startTime && thisTime - startTime < retryTime + 2,
              String.format(
                  "Waited %s seconds instead an expected %s seconds wait",
                  thisTime - startTime, retryTime));
        } catch (NoHostAvailableException e) {
          Thread.sleep(1000);
          continue;
        }
        break;
      }
      break;
    }
  }
}
