package com.netflix.runtime.health.core;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicInteger;

import com.netflix.runtime.health.api.Health;
import com.netflix.runtime.health.api.HealthIndicator;
import com.netflix.runtime.health.api.HealthIndicatorCallback;

/**
 */
public class SimpleHealthCheckAggregator implements HealthCheckAggregator {
  
	private final List<HealthIndicator> indicators;
	private final ScheduledExecutorService executor;
	private final TimeUnit units;
	private final long maxWaitTime;

    public SimpleHealthCheckAggregator(List<HealthIndicator> indicators, long maxWaitTime, TimeUnit units) {
        this.indicators = new ArrayList<>(indicators);
        this.maxWaitTime = maxWaitTime;
        this.units = units;
        this.executor = Executors.newSingleThreadScheduledExecutor();
    }

    public CompletableFuture<HealthCheckStatus> check() {
        final List<HealthIndicatorCallbackImpl> callbacks = new ArrayList<>(indicators.size());
        final CompletableFuture<HealthCheckStatus> future = new CompletableFuture<HealthCheckStatus>();
        final AtomicInteger counter = new AtomicInteger(indicators.size());
        
        indicators.stream().forEach(indicator -> {

            HealthIndicatorCallbackImpl callback = new HealthIndicatorCallbackImpl(indicator) {
                @Override
                public void inform(Health status) {
                    setHealth(status);
                    if (counter.decrementAndGet() == 0) {
                        future.complete(getStatusFromCallbacks(callbacks));
                    }
                }
            };

            callbacks.add(callback);
            indicator.check(callback);
        });
        if(indicators.size() == 0) {
        	future.complete(HealthCheckStatus.create(true, Collections.emptyList()));
        }
        
        if (maxWaitTime != 0 && units != null) {
            executor.schedule(new Runnable() {
                @Override
                public void run() {
                    future.complete(getStatusFromCallbacks(callbacks));
                }
            }, maxWaitTime, units);
        }
    
        return future;
    }

	protected HealthCheckStatus getStatusFromCallbacks(List<HealthIndicatorCallbackImpl> callbacks) {
	    List<Health> healths = new ArrayList<>(callbacks.size());
	    boolean isHealthy = callbacks.stream()
		    .map(callback -> {
		    	Health health = Health.from(callback.getHealthOrTimeout())
		    			.withDetail("className", callback.getIndicator().getClass().getName())
		    			.build();
		    	healths.add(health);
		    	return health;
		    })
		    .map(health -> health.isHealthy())
		    .reduce(true, (a,b) -> a && b); 
	    return HealthCheckStatus.create(isHealthy, healths);
	}
    
     
    abstract class HealthIndicatorCallbackImpl implements HealthIndicatorCallback {
        private volatile Health health;
        private final HealthIndicator indicator;
        
        HealthIndicatorCallbackImpl(HealthIndicator indicator) {
            this.indicator = indicator;
        }
        
        void setHealth(Health health) {
            this.health = health;
        }
        
        public Health getHealthOrTimeout() {
            return health != null 
                ? health
                : Health
                    .unhealthy(new TimeoutException("Timed out waiting for response"))
                    .build();
        }
        
        public HealthIndicator getIndicator() {
        	return this.indicator;
        }
    }
}
@@ -389,16 +389,52 @@
 Integer;
+%0Aimport java.util.stream.Collectors;
 %0A%0Aimport
@@ -1498,24 +1498,61 @@
     %0A       
+ List%3CCompletableFuture%3C?%3E%3E futures =
  indicators.
@@ -1564,15 +1564,11 @@
 m().
-forEach
+map
 (ind
@@ -2026,22 +2026,103 @@
 ;%0A      
+%0A            return CompletableFuture.runAsync(()-%3E %7B%0A%09            try %7B%0A%09   
       
+   %09
 indicato
@@ -2136,24 +2136,29 @@
 (callback);%0A
+%09    
         %7D);%0A
@@ -2154,18 +2154,238 @@
        %7D
-);
+%0A%09            catch(Exception ex) %0A%09            %7B%0A%09            %09callback.inform(Health.unhealthy(ex).build());%0A%09            %7D%0A            %7D);%0A            %0A        %7D).collect(Collectors.toList());%0A        %0A        %0A        
 %0A       
@@ -2670,24 +2670,24 @@
 oid run() %7B%0A
-
             
@@ -2738,32 +2738,146 @@
 ks(callbacks));%0A
+                    CompletableFuture.allOf(futures.toArray(new CompletableFuture%5Bfutures.size()%5D)).cancel(true);%0A
                 

stop errors from bubbling through aggregator make properly async add cancel logic on timeout