/*
 * Copyright (C) 2017-2017 DataStax Inc.
 *
 * Licensed 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.oss.driver.api.testinfra.utils;

import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.BooleanSupplier;

import static org.assertj.core.api.Fail.fail;

public class ConditionChecker {

  private static final int DEFAULT_PERIOD_MILLIS = 500;

  private static final int DEFAULT_TIMEOUT_MILLIS = 60000;

  public static class ConditionCheckerBuilder {

    private long timeout = DEFAULT_TIMEOUT_MILLIS;

    private TimeUnit timeoutUnit = TimeUnit.MILLISECONDS;

    private long period = DEFAULT_PERIOD_MILLIS;

    private TimeUnit periodUnit = TimeUnit.MILLISECONDS;

    private final BooleanSupplier predicate;

    private String description;

    ConditionCheckerBuilder(BooleanSupplier predicate) {
      this.predicate = predicate;
    }

    public ConditionCheckerBuilder every(long period, TimeUnit unit) {
      this.period = period;
      periodUnit = unit;
      return this;
    }

    public ConditionCheckerBuilder every(long periodMillis) {
      period = periodMillis;
      periodUnit = TimeUnit.MILLISECONDS;
      return this;
    }

    public ConditionCheckerBuilder before(long timeout, TimeUnit unit) {
      this.timeout = timeout;
      timeoutUnit = unit;
      return this;
    }

    public ConditionCheckerBuilder before(long timeoutMillis) {
      timeout = timeoutMillis;
      timeoutUnit = TimeUnit.MILLISECONDS;
      return this;
    }

    public ConditionCheckerBuilder as(String description) {
      this.description = description;
      return this;
    }

    public void becomesTrue() {
      new ConditionChecker(predicate, period, periodUnit, description).await(timeout, timeoutUnit);
    }

    public void becomesFalse() {
      new ConditionChecker(() -> !predicate.getAsBoolean(), period, periodUnit, description)
          .await(timeout, timeoutUnit);
    }
  }

  public static ConditionCheckerBuilder checkThat(BooleanSupplier predicate) {
    return new ConditionCheckerBuilder(predicate);
  }

  private final BooleanSupplier predicate;
  private final String description;
  private final Lock lock;
  private final Condition condition;
  private final Timer timer;

  public ConditionChecker(
      BooleanSupplier predicate, long period, TimeUnit periodUnit, String description) {
    this.predicate = predicate;
    this.description = (description != null) ? description : this.toString();
    lock = new ReentrantLock();
    condition = lock.newCondition();
    timer = new Timer("condition-checker", true);
    timer.schedule(
        new TimerTask() {
          @Override
          public void run() {
            checkCondition();
          }
        },
        0,
        periodUnit.toMillis(period));
  }

  /** Waits until the predicate becomes true, or a timeout occurs, whichever happens first. */
  public void await(long timeout, TimeUnit unit) {
    boolean interrupted = false;
    long nanos = unit.toNanos(timeout);
    lock.lock();
    try {
      while (!evalCondition()) {
        if (nanos <= 0L)
          fail(
              String.format(
                  "Timeout after %s %s while waiting for '%s'",
                  timeout, unit.toString().toLowerCase(), description));
        try {
          nanos = condition.awaitNanos(nanos);
        } catch (InterruptedException e) {
          interrupted = true;
        }
      }
    } finally {
      timer.cancel();
      if (interrupted) Thread.currentThread().interrupt();
    }
  }

  private void checkCondition() {
    lock.lock();
    try {
      if (evalCondition()) {
        condition.signal();
      }
    } finally {
      lock.unlock();
    }
  }

  private boolean evalCondition() {
    return predicate.getAsBoolean();
  }
}

@@ -1396,39 +1396,30 @@
 ivate final 
-BooleanSupplier
+Object
  predicate;%0A
@@ -1542,32 +1542,130 @@
 edicate;%0A    %7D%0A%0A
+    public ConditionCheckerBuilder(Runnable predicate) %7B%0A      this.predicate = predicate;%0A    %7D%0A%0A
     public Condi
@@ -2469,16 +2469,22 @@
 edicate,
+ true,
  period,
@@ -2508,16 +2508,27 @@
 ription)
+%0A          
 .await(t
@@ -2620,15 +2620,8 @@
 ker(
-() -%3E !
 pred
@@ -2625,31 +2625,23 @@
 redicate
-.getAsBoolean()
+, false
 , period
@@ -2860,45 +2860,205 @@
 %0A  p
-rivate final BooleanSupplier predicat
+ublic static ConditionCheckerBuilder checkThat(Runnable predicate) %7B%0A    return new ConditionCheckerBuilder(predicate);%0A  %7D%0A%0A  private final Object predicate;%0A  private final boolean expectedOutcom
 e;%0A 
@@ -3223,34 +3223,62 @@
     
-BooleanSupplier predicate,
+Object predicate,%0A      boolean expectedOutcome,%0A     
  lon
@@ -3278,32 +3278,38 @@
     long period,
+%0A     
  TimeUnit period
@@ -3313,16 +3313,22 @@
 iodUnit,
+%0A     
  String 
@@ -3370,24 +3370,68 @@
  predicate;%0A
+    this.expectedOutcome = expectedOutcome;%0A
     this.des
@@ -4778,22 +4778,91 @@
 ) %7B%0A    
-return
+if (predicate instanceof BooleanSupplier) %7B%0A      return ((BooleanSupplier)
  predica
@@ -4863,16 +4863,17 @@
 redicate
+)
 .getAsBo
@@ -4879,16 +4879,378 @@
 oolean()
-;
+ == expectedOutcome;%0A    %7D else if (predicate instanceof Runnable) %7B%0A      boolean succeeded = true;%0A      try %7B%0A        ((Runnable) predicate).run();%0A      %7D catch (Throwable t) %7B%0A        succeeded = false;%0A      %7D%0A      return succeeded == expectedOutcome;%0A    %7D else %7B%0A      throw new AssertionError(%22Unsupported predicate type %22 + predicate.getClass());%0A    %7D
 %0A  %7D%0A%7D%0A

Allow usage of a Runnable in ConditionChecker