package com.airlock.iam.oauth2.domain.service.metadata;

import com.airlock.iam.base.api.domain.model.exception.FailureException;
import com.airlock.iam.base.misc.log.LogUtil;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.oauth2.domain.model.Issuer;
import com.airlock.iam.oauth2.domain.model.OAuth2ResponseMode;
import com.airlock.iam.oauth2.domain.model.PkceCodeChallengeMethod;
import com.airlock.iam.oauth2.domain.model.acr.OpenIdConnectAcrValue;
import com.airlock.iam.oauth2.domain.model.authorization.OAuth2ResponseType;
import com.airlock.iam.oauth2.domain.model.claims.ClaimType;
import com.airlock.iam.oauth2.domain.model.metadata.OAuth2EndpointUriComposer;
import com.airlock.iam.oauth2.domain.model.metadata.OAuth2MetadataIssuerValidationException;
import com.airlock.iam.oauth2.domain.model.metadata.OAuth2MetadataSettings;
import com.airlock.iam.oauth2.domain.model.metadata.UILocale;
import com.google.common.base.Preconditions;
import com.nimbusds.langtag.LangTag;
import com.nimbusds.langtag.LangTagException;
import com.nimbusds.oauth2.sdk.GrantType;
import com.nimbusds.oauth2.sdk.ParseException;
import com.nimbusds.oauth2.sdk.ResponseMode;
import com.nimbusds.oauth2.sdk.ResponseType;
import com.nimbusds.oauth2.sdk.Scope;
import com.nimbusds.oauth2.sdk.Scope.Value;
import com.nimbusds.oauth2.sdk.as.AuthorizationServerMetadata;
import com.nimbusds.oauth2.sdk.auth.ClientAuthenticationMethod;
import com.nimbusds.oauth2.sdk.pkce.CodeChallengeMethod;
import com.nimbusds.openid.connect.sdk.OIDCScopeValue;
import com.nimbusds.openid.connect.sdk.SubjectType;
import com.nimbusds.openid.connect.sdk.claims.ACR;
import com.nimbusds.openid.connect.sdk.op.OIDCProviderMetadata;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public final class OAuth2MetadataExtractor {
   private static final ELogger LOG = ELoggerFactory.getLogger(OAuth2MetadataExtractor.class);
   private final OAuth2EndpointUriComposer uriComposer;
   private final OAuth2MetadataIssuerExtractor issuerExtractor;

   public OAuth2MetadataExtractor(OAuth2EndpointUriComposer uriComposer) {
      this.uriComposer = uriComposer;
      this.issuerExtractor = new OAuth2MetadataIssuerExtractor();
   }

   public OIDCProviderMetadata extractOIDCProviderMetadata(OAuth2MetadataSettings settings) throws OAuth2MetadataIssuerValidationException {
      Preconditions.checkState(
         settings.openIDConnectSettings().isPresent(), "Cannot generate discovery endpoint metadata for AS that does not support OpenID Connect"
      );
      Issuer issuer = this.determineIssuer(settings);
      OIDCProviderMetadata result = new OIDCProviderMetadata(mapIssuer(issuer), List.of(SubjectType.PUBLIC), this.uriComposer.jwksURI());
      result.setIDTokenJWSAlgs(settings.openIDConnectSettings().get().idTokenSigningAlgorithms());
      List<OpenIdConnectAcrValue> openIdConnectAcrValues = settings.openIDConnectSettings().get().acrValuesSupported();
      result.setACRs(nullIfEmpty(mapACRValues(openIdConnectAcrValues)));
      result.setClaims(settings.openIDConnectSettings().get().claimsSupported());
      List<ClaimType> claimTypes = settings.openIDConnectSettings().get().claimTypesSupported();
      result.setClaimTypes(mapClaimTypes(claimTypes));
      result.setCheckSessionIframeURI(settings.checkSessionSupported() ? this.uriComposer.checkSessionIframeURI() : null);
      settings.userInfoEndpointSettings().ifPresent(endpointSettings -> this.fillUserInfoEndpointSetting(result));
      this.fillBackChannelLogoutMetadata(result, settings);
      this.fillGeneralMetadata(result, settings);
      return result;
   }

   private static <T> List<T> nullIfEmpty(List<T> list) {
      return list != null && !list.isEmpty() ? list : null;
   }

   private static List<ACR> mapACRValues(List<OpenIdConnectAcrValue> openIdConnectAcrValues) {
      return openIdConnectAcrValues.stream().map(openIdConnectAcrValue -> new ACR(openIdConnectAcrValue.name())).toList();
   }

   private static List<com.nimbusds.openid.connect.sdk.claims.ClaimType> mapClaimTypes(List<ClaimType> claimTypes) {
      return claimTypes.stream().map(claimType -> com.nimbusds.openid.connect.sdk.claims.ClaimType.valueOf(claimType.name())).toList();
   }

   public AuthorizationServerMetadata extractAuthorizationServerMetadata(OAuth2MetadataSettings settings) throws OAuth2MetadataIssuerValidationException {
      if (settings.openIDConnectSettings().isPresent()) {
         return this.extractOIDCProviderMetadata(settings);
      }

      Issuer issuer = this.determineIssuer(settings);
      AuthorizationServerMetadata result = new AuthorizationServerMetadata(mapIssuer(issuer));
      this.fillGeneralMetadata(result, settings);
      return result;
   }

   private Issuer determineIssuer(OAuth2MetadataSettings settings) throws OAuth2MetadataIssuerValidationException {
      return this.issuerExtractor
         .determineIssuer(this.uriComposer.providerBaseURI(), Issuer.issuer(settings.issuer().getValue()), settings.issuerValidationMode());
   }

   private static com.nimbusds.oauth2.sdk.id.Issuer mapIssuer(Issuer issuer) {
      return new com.nimbusds.oauth2.sdk.id.Issuer(issuer.getValue());
   }

   private void fillGeneralMetadata(AuthorizationServerMetadata result, OAuth2MetadataSettings settings) {
      settings.registrationEndpointSettings().ifPresent(endpointSettings -> this.fillRegistrationEndpointMetadata(result));
      settings.authorizationEndpointSettings().ifPresent(endpointSettings -> this.fillAuthorizationEndpointMetadata(result));
      settings.tokenEndpointSettings().ifPresent(endpointSettings -> this.fillTokenEndpointMetadata(result, endpointSettings));
      settings.pushedAuthorizationRequestEndpointSettings().ifPresent(endpointSettings -> this.fillParEndpointMetadata(result, endpointSettings));
      settings.jwksEndpointSettings().ifPresent(endpointSettings -> this.fillJwksEndpointMetadata(result));
      settings.tokenRevocationEndpointSettings().ifPresent(endpointSettings -> this.fillRevocationEndpointMetadata(result, endpointSettings));
      settings.tokenIntrospectionEndpointSettings().ifPresent(endpointSettings -> this.fillIntrospectionEndpointMetadata(result, endpointSettings));
      determineScope(settings).ifPresent(result::setScopes);
      result.setResponseTypes(determineResponseTypes(settings));
      result.setResponseModes(determineResponseModes(settings));
      result.setGrantTypes(determineGrantTypes(settings));
      determineCodeChallengeMethods(settings).ifPresent(result::setCodeChallengeMethods);
      result.setSupportsRequestURIParam(false);
      result.setSupportsRequestParam(false);
      this.mapUILocales(settings.uiLocales()).ifPresent(result::setUILocales);
   }

   private Optional<List<LangTag>> mapUILocales(List<UILocale> uiLocales) {
      List<LangTag> langTags = uiLocales.stream().flatMap(l -> {
         try {
            return Stream.of(LangTag.parse(l.getValue()));
         } catch (LangTagException e) {
            LOG.info(String.format("Failed to add %s as ui_locale", LogUtil.quote(l.getValue())), e);
            return Stream.empty();
         }
      }).collect(Collectors.toList());
      return langTags.isEmpty() ? Optional.empty() : Optional.of(langTags);
   }

   private void fillAuthorizationEndpointMetadata(AuthorizationServerMetadata result) {
      result.setAuthorizationEndpointURI(this.uriComposer.authorizationEndpointURI());
   }

   private void fillJwksEndpointMetadata(AuthorizationServerMetadata metadata) {
      metadata.setJWKSetURI(this.uriComposer.jwksURI());
   }

   private void fillRegistrationEndpointMetadata(AuthorizationServerMetadata metadata) {
      metadata.setRegistrationEndpointURI(this.uriComposer.registrationEndpointURI());
   }

   private void fillTokenEndpointMetadata(AuthorizationServerMetadata metadata, OAuth2MetadataSettings.TokenEndpointSettings endpointSettings) {
      metadata.setTokenEndpointURI(this.uriComposer.tokenEndpointURI());
      metadata.setTokenEndpointAuthMethods(
         endpointSettings.supportedAuthenticationMethods().stream().map(method -> new ClientAuthenticationMethod(method.getIdentifier())).toList()
      );
   }

   private void fillParEndpointMetadata(
      AuthorizationServerMetadata metadata, OAuth2MetadataSettings.PushedAuthorizationRequestEndpointSettings endpointSettings
   ) {
      metadata.setPushedAuthorizationRequestEndpointURI(this.uriComposer.pushedAuthorizationRequestURI());
      metadata.requiresPushedAuthorizationRequests(endpointSettings.pushedAuthorizationRequestsRequired());
   }

   private void fillRevocationEndpointMetadata(AuthorizationServerMetadata metadata, OAuth2MetadataSettings.TokenRevocationEndpointSettings endpointSettings) {
      metadata.setRevocationEndpointURI(this.uriComposer.revocationEndpointURI());
      metadata.setRevocationEndpointAuthMethods(List.of(new ClientAuthenticationMethod(endpointSettings.authenticationMethod().getIdentifier())));
   }

   private void fillIntrospectionEndpointMetadata(
      AuthorizationServerMetadata metadata, OAuth2MetadataSettings.TokenIntrospectionEndpointSettings endpointSettings
   ) {
      metadata.setIntrospectionEndpointURI(this.uriComposer.introspectionEndpointURI());
      metadata.setIntrospectionEndpointAuthMethods(List.of(new ClientAuthenticationMethod(endpointSettings.authenticationMethod().getIdentifier())));
   }

   private void fillUserInfoEndpointSetting(OIDCProviderMetadata metadata) {
      metadata.setUserInfoEndpointURI(this.uriComposer.userInfoEndpointURI());
   }

   private void fillBackChannelLogoutMetadata(OIDCProviderMetadata result, OAuth2MetadataSettings settings) {
      if (settings.isBackChannelLogoutSupported()) {
         result.setSupportsBackChannelLogout(true);
         result.setSupportsBackChannelLogoutSession(true);
      }
   }

   private static Optional<Scope> determineScope(OAuth2MetadataSettings settings) {
      return settings.openIDConnectSettings().isPresent() ? Optional.of(new Scope(new Value[]{OIDCScopeValue.OPENID})) : Optional.empty();
   }

   private static List<ResponseType> determineResponseTypes(OAuth2MetadataSettings settings) {
      List<ResponseType> responseTypes = new ArrayList<>();
      if (settings.authorizationCodeGrantSupported()) {
         responseTypes.add(ResponseType.CODE);
      }

      settings.openIDConnectSettings()
         .map(OAuth2MetadataSettings.OpenIDConnectSettings::hybridFlowResponseTypesSupported)
         .ifPresent(oAuth2ResponseTypes -> responseTypes.addAll(convertResponseTypes((Set<OAuth2ResponseType>)oAuth2ResponseTypes)));
      return responseTypes;
   }

   private static List<ResponseType> convertResponseTypes(Set<OAuth2ResponseType> oAuth2ResponseTypes) {
      return oAuth2ResponseTypes.stream().map(OAuth2MetadataExtractor::convertResponseType).collect(Collectors.toList());
   }

   private static ResponseType convertResponseType(OAuth2ResponseType oAuth2ResponseType) {
      try {
         return ResponseType.parse(oAuth2ResponseType.responseType());
      } catch (ParseException e) {
         throw new FailureException(String.format("Could not parse OAuth 2.0 response type: '%s'", oAuth2ResponseType.responseType()), e);
      }
   }

   private static List<ResponseMode> determineResponseModes(OAuth2MetadataSettings settings) {
      return Stream.concat(
            settings.supportedResponseModes().getAuthCodeFlowResponseModes().stream(), settings.supportedResponseModes().getHybridFlowResponseModes().stream()
         )
         .map(OAuth2ResponseMode::getValue)
         .distinct()
         .<ResponseMode>map(ResponseMode::new)
         .toList();
   }

   private static List<GrantType> determineGrantTypes(OAuth2MetadataSettings settings) {
      List<GrantType> result = new ArrayList<>();
      if (settings.authorizationCodeGrantSupported()) {
         result.add(GrantType.AUTHORIZATION_CODE);
      }

      if (settings.clientCredentialsGrantSupported()) {
         result.add(GrantType.CLIENT_CREDENTIALS);
      }

      if (settings.tokenExchangeGrantSupported()) {
         result.add(GrantType.TOKEN_EXCHANGE);
      }

      if (settings.generateRefreshToken()) {
         result.add(GrantType.REFRESH_TOKEN);
      }

      return result;
   }

   private static Optional<List<CodeChallengeMethod>> determineCodeChallengeMethods(OAuth2MetadataSettings endpointSettings) {
      Optional<PkceCodeChallengeMethod> pkceCodeChallengeMethod = endpointSettings.pkceCodeChallengeMethod();
      if (pkceCodeChallengeMethod.isEmpty()) {
         return Optional.empty();
      } else {
         return pkceCodeChallengeMethod.get() == PkceCodeChallengeMethod.S256
            ? Optional.of(List.of(CodeChallengeMethod.S256))
            : Optional.of(List.of(CodeChallengeMethod.PLAIN, CodeChallengeMethod.S256));
      }
   }
}