package io.quarkus.oidc.runtime;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.function.Function;

import org.jboss.logging.Logger;

import io.quarkus.arc.runtime.BeanContainer;
import io.quarkus.oidc.OIDCException;
import io.quarkus.runtime.RuntimeValue;
import io.quarkus.runtime.annotations.Recorder;
import io.vertx.core.AsyncResult;
import io.vertx.core.Handler;
import io.vertx.core.Vertx;
import io.vertx.ext.auth.PubSecKeyOptions;
import io.vertx.ext.auth.oauth2.OAuth2Auth;
import io.vertx.ext.auth.oauth2.OAuth2ClientOptions;
import io.vertx.ext.auth.oauth2.providers.KeycloakAuth;
import io.vertx.ext.jwt.JWTOptions;

@Recorder
public class OidcRecorder {

    private static final Logger LOG = Logger.getLogger(OidcRecorder.class);

    public void setup(OidcConfig config, RuntimeValue<Vertx> vertx, BeanContainer beanContainer) {
        final Vertx vertxValue = vertx.getValue();
        Map<String, TenantConfigContext> tenantsConfig = new HashMap<>();

        for (Map.Entry<String, OidcTenantConfig> tenant : config.namedTenants.entrySet()) {
            if (config.defaultTenant.getTenantId().isPresent()
                    && tenant.getKey().equals(config.defaultTenant.getTenantId().get())) {
                throw new OIDCException("tenant-id '" + tenant.getKey() + "' duplicates the default tenant-id");
            }
            if (tenant.getValue().getTenantId().isPresent() && !tenant.getKey().equals(tenant.getValue().getTenantId().get())) {
                throw new OIDCException("Configuration has 2 different tenant-id values: '"
                        + tenant.getKey() + "' and '" + tenant.getValue().getTenantId().get() + "'");
            }
            tenantsConfig.put(tenant.getKey(), createTenantContext(vertxValue, tenant.getValue()));
        }

        DefaultTenantConfigResolver resolver = beanContainer.instance(DefaultTenantConfigResolver.class);

        resolver.setDefaultTenant(createTenantContext(vertxValue, config.defaultTenant));
        resolver.setTenantsConfig(tenantsConfig);
        resolver.setTenantConfigContextFactory(new Function<OidcTenantConfig, TenantConfigContext>() {
            @Override
            public TenantConfigContext apply(OidcTenantConfig config) {
                return createTenantContext(vertxValue, config);
            }
        });
    }

    private TenantConfigContext createTenantContext(Vertx vertx, OidcTenantConfig oidcConfig) {
        OAuth2ClientOptions options = new OAuth2ClientOptions();

        if (!oidcConfig.getAuthServerUrl().isPresent()) {
            return null;
        }

        // Base IDP server URL
        options.setSite(oidcConfig.getAuthServerUrl().get());
        // RFC7662 introspection service address
        if (oidcConfig.getIntrospectionPath().isPresent()) {
            options.setIntrospectionPath(oidcConfig.getIntrospectionPath().get());
        }

        // RFC7662 JWKS service address
        if (oidcConfig.getJwksPath().isPresent()) {
            options.setJwkPath(oidcConfig.getJwksPath().get());
        }

        if (oidcConfig.getClientId().isPresent()) {
            options.setClientID(oidcConfig.getClientId().get());
        }

        if (oidcConfig.getCredentials().secret.isPresent()) {
            options.setClientSecret(oidcConfig.getCredentials().secret.get());
        }
        if (oidcConfig.getPublicKey().isPresent()) {
            options.addPubSecKey(new PubSecKeyOptions()
                    .setAlgorithm("RS256")
                    .setPublicKey(oidcConfig.getPublicKey().get()));
        }
        if (oidcConfig.getToken().issuer.isPresent()) {
            options.setValidateIssuer(false);
        }

        if (oidcConfig.getToken().getExpirationGrace().isPresent()) {
            JWTOptions jwtOptions = new JWTOptions();
            jwtOptions.setLeeway(oidcConfig.getToken().getExpirationGrace().get());
            options.setJWTOptions(jwtOptions);
        }

        final long connectionDelayInSecs = oidcConfig.getConnectionDelay().isPresent()
                ? oidcConfig.getConnectionDelay().get().toMillis() / 1000
                : 0;
        final long connectionRetryCount = connectionDelayInSecs > 1 ? connectionDelayInSecs / 2 : 1;
        if (connectionRetryCount > 1) {
            LOG.infof("Connecting to IDP for up to %d times every 2 seconds", connectionRetryCount);
        }

        OAuth2Auth auth = null;
        for (long i = 0; i < connectionRetryCount; i++) {
            try {
                CompletableFuture<OAuth2Auth> cf = new CompletableFuture<>();
                KeycloakAuth.discover(vertx, options, new Handler<AsyncResult<OAuth2Auth>>() {
                    @Override
                    public void handle(AsyncResult<OAuth2Auth> event) {
                        if (event.failed()) {
                            cf.completeExceptionally(toOidcException(event.cause()));
                        } else {
                            cf.complete(event.result());
                        }
                    }
                });

                auth = cf.join();
                break;
            } catch (OIDCException ex) {
                if (i + 1 < connectionRetryCount) {
                    try {
                        Thread.sleep(2000);
                    } catch (InterruptedException iex) {
                        // continue connecting
                    }

                } else {
                    throw ex;
                }
            }
        }

        return new TenantConfigContext(auth, oidcConfig);
    }

    protected static OIDCException toOidcException(Throwable cause) {
        final String message = "OIDC server is not available at the 'quarkus.oidc.auth-server-url' URL. "
                + "Please make sure it is correct. Note it has to end with a realm value if you work with Keycloak, for example:"
                + " 'https://localhost:8180/auth/realms/quarkus'";
        return new OIDCException(message, cause);
    }

}

@@ -122,16 +122,65 @@
 Future;%0A
+import java.util.concurrent.CompletionException;%0A
 import j
@@ -5241,28 +5241,264 @@
 ch (
-OIDCException ex) %7B%0A
+Throwable throwable) %7B%0A                while (throwable instanceof CompletionException && throwable.getCause() != null) %7B%0A                    throwable = throwable.getCause();%0A                %7D%0A                if (throwable instanceof OIDCException) %7B%0A    
     
@@ -5557,38 +5557,46 @@
                 
+  
+  
 try %7B%0A
+    
                 
@@ -5635,32 +5635,36 @@
                 
+    
 %7D catch (Interru
@@ -5708,16 +5708,20 @@
         
+ 
+   
 // conti
@@ -5735,16 +5735,128 @@
 necting%0A
+                        %7D%0A                    %7D else %7B%0A                        throw (OIDCException) throwable;%0A
         
@@ -5861,33 +5861,32 @@
               %7D%0A
-%0A
                 
@@ -5920,18 +5920,44 @@
   throw 
-ex
+new OIDCException(throwable)
 ;%0A      

Fix a connection-delay processing