package io.quarkus.consul.config.runtime;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletionException;
import java.util.function.Consumer;

import org.eclipse.microprofile.config.spi.ConfigSource;
import org.eclipse.microprofile.config.spi.ConfigSourceProvider;
import org.jboss.logging.Logger;

import io.smallrye.mutiny.Uni;

class ConsulConfigSourceProvider implements ConfigSourceProvider {

    private static final Logger log = Logger.getLogger(ConsulConfigSourceProvider.class);

    private final ConsulConfig config;

    private final ConsulConfigGateway consulConfigGateway;
    private final ResponseConfigSourceUtil responseConfigSourceUtil;

    public ConsulConfigSourceProvider(ConsulConfig config) {
        this(config, new VertxConsulConfigGateway(config), new ResponseConfigSourceUtil());
    }

    // visible for testing
    ConsulConfigSourceProvider(ConsulConfig config, ConsulConfigGateway consulConfigGateway) {
        this(config, consulConfigGateway, new ResponseConfigSourceUtil());
    }

    private ConsulConfigSourceProvider(ConsulConfig config, ConsulConfigGateway consulConfigGateway,
            ResponseConfigSourceUtil responseConfigSourceUtil) {
        this.config = config;
        this.consulConfigGateway = consulConfigGateway;
        this.responseConfigSourceUtil = responseConfigSourceUtil;
    }

    @Override
    public Iterable<ConfigSource> getConfigSources(ClassLoader cl) {
        Map<String, ValueType> keys = config.keysAsMap();
        if (keys.isEmpty()) {
            log.debug("No keys were configured for config source lookup");
            return Collections.emptyList();
        }

        List<ConfigSource> result = new ArrayList<>(keys.size());

        List<Uni<?>> allUnis = new ArrayList<>();

        for (Map.Entry<String, ValueType> entry : keys.entrySet()) {
            String fullKey = config.prefix.isPresent() ? config.prefix.get() + "/" + entry.getKey() : entry.getKey();
            allUnis.add(consulConfigGateway.getValue(fullKey).invoke(new Consumer<Response>() {
                @Override
                public void accept(Response response) {
                    if (response != null) {
                        result.add(
                                responseConfigSourceUtil.toConfigSource(response, entry.getValue(),
                                        config.prefix));
                    } else {
                        String message = "Key '" + fullKey + "' not found in Consul.";
                        if (config.failOnMissingKey) {
                            throw new RuntimeException(message);
                        } else {
                            log.info(message);
                        }
                    }
                }
            }));
        }

        try {
            Uni.combine().all().unis(allUnis).discardItems().await()
                    .atMost(config.agent.connectionTimeout.plus(config.agent.readTimeout.multipliedBy(2)));
        } catch (CompletionException e) {
            throw new RuntimeException("An error occurred while attempting to fetch configuration from Consul.", e);
        } finally {
            consulConfigGateway.close();
        }

        return result;
    }
}

@@ -411,16 +411,59 @@
 iny.Uni;
+%0Aimport io.smallrye.mutiny.groups.UniAwait;
 %0A%0Aclass 
@@ -2929,16 +2929,39 @@
         
+ UniAwait%3CVoid%3E await =
  Uni.com
@@ -3009,16 +3009,17 @@
 .await()
+;
 %0A       
@@ -3023,24 +3023,180 @@
         
-        
+if (config.agent.connectionTimeout.isZero() && config.agent.readTimeout.isZero()) %7B%0A                await.indefinitely();%0A            %7D else %7B%0A                await
 .atMost(
@@ -3271,24 +3271,38 @@
 iedBy(2)));%0A
+            %7D%0A
         %7D ca

Fix potential IllegalArgumentException