package org.springframework.cloud.zookeeper;

import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.springframework.boot.actuate.endpoint.Endpoint;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * @author Spencer Gibb
 */
@Configuration
@ConditionalOnProperty(value = "zookeeper.enabled", matchIfMissing = true)
@EnableConfigurationProperties
public class ZookeeperAutoConfiguration {

	@Bean
	@ConditionalOnMissingBean
	public ZookeeperProperties zookeeperProperties() {
		return new ZookeeperProperties();
	}

	@Bean(destroyMethod = "close")
	@ConditionalOnMissingBean
	public CuratorFramework curatorFramework(RetryPolicy retryPolicy) {
		CuratorFramework curator = CuratorFrameworkFactory.builder()
				.retryPolicy(retryPolicy)
				// TODO: support ensembleProvider via ExhibitorEnsembleProvider
				// .ensembleProvider(new ExhibitorEnsembleProvider())
				.connectString(zookeeperProperties().getConnectString()).build();
		curator.start();
		return curator;
	}

	@Bean
	@ConditionalOnMissingBean
	public RetryPolicy exponentialBackoffRetry() {
		return new ExponentialBackoffRetry(zookeeperProperties().getBaseSleepTimeMs(),
				zookeeperProperties().getMaxRetries(),
				zookeeperProperties().getMaxSleepMs());
	}

	@Configuration
	@ConditionalOnClass(Endpoint.class)
	protected static class ZookeeperHealthConfig {
		@Bean
		@ConditionalOnMissingBean
		public ZookeeperEndpoint zookeeperEndpoint() {
			return new ZookeeperEndpoint();
		}
        @Bean
        @ConditionalOnMissingBean
        public ZookeeperHealthIndicator zookeeperHealthIndicator() {
		return new ZookeeperHealthIndicator();
	}
    }
}

@@ -78,16 +78,69 @@
 Policy;%0A
+import org.apache.curator.ensemble.EnsembleProvider;%0A
 import o
@@ -299,24 +299,87 @@
 ckoffRetry;%0A
+import org.springframework.beans.factory.annotation.Autowired;%0A
 import org.s
@@ -1054,16 +1054,91 @@
 tion %7B%0A%0A
+%09@Autowired(required = false)%0A%09private EnsembleProvider ensembleProvider;%0A%0A
 %09@Bean%0A%09
@@ -1401,23 +1401,38 @@
 ramework
- curato
+Factory.Builder builde
 r = Cura
@@ -1464,91 +1464,56 @@
 er()
+;
 %0A%09%09
-%09%09.retryPolicy(retryPolicy)%0A%09%09%09%09// TODO: support ensembleProvider via ExhibitorE
+if (ensembleProvider != null) %7B%0A%09%09%09builder.e
 nsem
@@ -1527,17 +1527,9 @@
 ider
-%0A%09%09%09%09// .
+(
 ense
@@ -1544,40 +1544,80 @@
 ider
-(new ExhibitorEnsembleProvider()
+);%0A%09%09%7D%0A%09%09CuratorFramework curator = builder%0A%09%09%09%09.retryPolicy(retryPolicy
 )%0A%09%09
@@ -1674,16 +1674,21 @@
 tring())
+%0A%09%09%09%09
 .build()

add support for optional EnsembleProvider