package com.airlock.iam.oauth2.domain.model.client.communication;

import com.airlock.iam.base.misc.log.LogUtil;
import com.airlock.iam.common.api.domain.model.claims.ClaimsValidationResult;
import com.airlock.iam.common.api.domain.model.claims.ClaimsValidator;
import com.airlock.iam.common.domain.model.claims.ClaimsValidationResultImpl;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.oauth2.domain.model.client.AudienceClaimValidationMethod;
import com.airlock.iam.oauth2.domain.model.client.OAuth2ClientSession;
import com.airlock.iam.oauth2.domain.model.client.OAuth2ClientSideCredentials;
import com.airlock.iam.oauth2.domain.model.client.OpenIdConnectClaim;
import com.airlock.iam.oauth2.domain.model.client.OpenIdConnectParam;
import com.airlock.iam.oauth2.domain.model.logging.OAuth2LogTypes;
import com.airlock.iam.oauth2.domain.model.token.IdToken;
import com.nimbusds.oauth2.sdk.id.Identifier;
import com.nimbusds.openid.connect.sdk.Nonce;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.Collection;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import lombok.NonNull;
import org.apache.commons.lang3.StringUtils;

public class OpenIDConnectClientParserImpl extends OAuth2ClientParserImpl {
   @NonNull
   private final Supplier<OpenIDConnectIdTokenSignatureValidator> signatureValidator;
   @NonNull
   private final List<ClaimsValidator> additionalClaimValidators;
   @NonNull
   private final OpenIDConnectClientParserSettings settings;

   public OpenIDConnectClientParserImpl(
      @NonNull ELogger logger,
      @NonNull Supplier<OpenIDConnectIdTokenSignatureValidator> signatureValidator,
      @NonNull List<ClaimsValidator> additionalClaimValidators,
      @NonNull OpenIDConnectClientParserSettings settings
   ) {
      super(logger);
      if (logger == null) {
         throw new NullPointerException("logger is marked non-null but is null");
      }

      if (signatureValidator == null) {
         throw new NullPointerException("signatureValidator is marked non-null but is null");
      }

      if (additionalClaimValidators == null) {
         throw new NullPointerException("additionalClaimValidators is marked non-null but is null");
      }

      if (settings == null) {
         throw new NullPointerException("settings is marked non-null but is null");
      }

      this.signatureValidator = signatureValidator;
      this.additionalClaimValidators = additionalClaimValidators;
      this.settings = settings;
   }

   public boolean validateIdToken(
      Map<String, Object> parsedAccessTokenResponse, OAuth2ClientSession clientSession, OAuth2ClientSideCredentials clientCredentials
   ) {
      this.getLogger().debug("Validating OpenID Connect ID Token");
      Optional<IdToken> idToken = this.extractIdToken(parsedAccessTokenResponse);
      if (idToken.isEmpty()) {
         this.getLogger().debug("OpenID Connect ID Token is missing - validation fails");
         return false;
      } else {
         return this.isValidIdToken(idToken.get(), clientSession, clientCredentials);
      }
   }

   public Optional<IdToken> extractIdToken(Map<String, Object> parsedAccessTokenResponse) {
      try {
         String idToken;
         if (parsedAccessTokenResponse.containsKey(OpenIdConnectParam.ID_TOKEN.getParameterName())) {
            idToken = (String)parsedAccessTokenResponse.get(OpenIdConnectParam.ID_TOKEN.getParameterName());
         } else {
            if (!parsedAccessTokenResponse.containsKey(OpenIdConnectParam.AUTHENTICATION_TOKEN.getParameterName())) {
               this.getLogger().info("Error retrieving OpenID Connect ID Token: ID Token not found!");
               return Optional.empty();
            }

            idToken = (String)parsedAccessTokenResponse.get(OpenIdConnectParam.AUTHENTICATION_TOKEN.getParameterName());
         }

         if (StringUtils.isNotBlank(idToken)) {
            return IdToken.parse(idToken, this.getLogger());
         }

         this.getLogger().info("Error retrieving OpenID Connect ID Token: ID Token is blank!");
      } catch (ClassCastException e) {
         this.getLogger().warn(OAuth2LogTypes.OAUTH2_MESSAGE, "Error retrieving OpenID Connect ID Token: unexpected format: " + e.getMessage());
      }

      return Optional.empty();
   }

   private boolean isValidIdToken(IdToken idToken, OAuth2ClientSession clientSession, OAuth2ClientSideCredentials clientCredentials) {
      boolean validSignature = this.signatureValidator.get().validateIdTokenSignature(idToken, clientCredentials);
      return validSignature && this.validateIdTokenClaims(idToken, clientSession);
   }

   public Map<String, Object> getParsedIdTokenPayLoad(Map<String, Object> parsedAccessTokenResponse) {
      Optional<IdToken> idToken = this.extractIdToken(parsedAccessTokenResponse);
      return idToken.isEmpty() ? new HashMap<>() : idToken.get().claims();
   }

   private boolean validateIdTokenClaims(IdToken idToken, OAuth2ClientSession clientSession) {
      Map<String, Object> claims = idToken.claims();
      if (!this.validateIdTokenAudience(claims)) {
         return false;
      } else if (!this.validateIdTokenIssuer(claims)) {
         return false;
      } else if (!this.validateAdditionalClaims(claims)) {
         return false;
      } else if (!this.validateAcrClaims(claims)) {
         return false;
      } else {
         return !this.validateNonce(claims, clientSession) ? false : this.validateIdTokenNotExpired(claims);
      }
   }

   private boolean validateNonce(Map<String, Object> claims, OAuth2ClientSession clientSession) {
      Object nonceClaim = claims.get(OpenIdConnectClaim.NONCE.getClaimName());
      String nonceValue = Objects.toString(nonceClaim, null);
      Optional<Nonce> nonce = Optional.ofNullable(Nonce.parse(nonceValue));
      if (nonce.isEmpty() && !this.settings.includeNonce()) {
         return true;
      } else {
         boolean isValidNonce = clientSession.hasNonce(nonce.<String>map(Identifier::getValue).orElse(null));
         if (!isValidNonce) {
            this.getLogger().info("OpenID Connect ID Token validation failed due to invalid nonce.");
            this.getLogger()
               .debug("Received ID Token nonce: " + nonce.<String>map(Identifier::getValue).orElse("<none>") + ". Expected: " + clientSession.getNonce());
            return false;
         } else {
            this.getLogger().debug("OpenID Connect ID Token contains valid nonce: " + nonce.get().getValue());
            return true;
         }
      }
   }

   private boolean validateAdditionalClaims(Map<String, Object> claims) {
      List<ClaimsValidationResult> invalidClaims = this.additionalClaimValidators
         .stream()
         .map(validator -> validator.validate(claims))
         .filter(res -> !res.isValid())
         .collect(Collectors.toList());
      if (invalidClaims.isEmpty()) {
         this.getLogger().debug("OpenID Connect ID Token passes custom claim validation");
         return true;
      } else {
         this.getLogger()
            .info(
               "OpenID Connect ID Token contains claims that do not pass custom validation: " + ClaimsValidationResultImpl.invalidClaimsAsString(invalidClaims)
            );
         return false;
      }
   }

   private boolean validateAcrClaims(Map<String, Object> claims) {
      if (!this.settings.enableAcrValidation()) {
         return true;
      }

      Object idTokenAcrClaim = claims.get(OpenIdConnectClaim.ACR.getClaimName());
      if (!this.validateAcrValues(idTokenAcrClaim)) {
         return false;
      }

      if (!this.validateClaimsRequest(idTokenAcrClaim)) {
         return false;
      }

      this.getLogger().debug("OpenID Connect ID Token contains one of the requested acr claims: " + idTokenAcrClaim);
      return true;
   }

   private boolean validateAcrValues(Object idTokenAcrClaim) {
      if (this.settings.acrValuesClaim().isEmpty()) {
         return true;
      } else if (idTokenAcrClaim == null) {
         this.infoLogMissingAcrClaim(this.settings.acrValuesClaim());
         return false;
      } else if (!this.settings.acrValuesClaim().contains(idTokenAcrClaim.toString())) {
         this.infoLogNonMatchingAcrClaim(this.settings.acrValuesClaim(), idTokenAcrClaim);
         return false;
      } else {
         return true;
      }
   }

   private boolean validateClaimsRequest(Object idTokenAcrClaim) {
      Optional<ClaimRequestSettings> acrClaim = this.settings
         .claimRequestSettings()
         .stream()
         .filter(claimRequestSettings -> ClaimRequestSettings.ClaimRequestMember.ID_TOKEN.equals(claimRequestSettings.topLevelMember()))
         .filter(c -> c.claim().equals(OpenIdConnectClaim.ACR.getClaimName()))
         .findFirst();
      if (acrClaim.isEmpty()) {
         return true;
      }

      List<String> requestedAcrs = acrClaim.get().values();
      if (idTokenAcrClaim == null) {
         if (requestedAcrs.isEmpty()) {
            this.getLogger().info("OpenID Connect ID Token validation failed. No acr claim is present but the client requested it");
         } else {
            this.infoLogMissingAcrClaim(requestedAcrs);
         }

         return false;
      } else if (!requestedAcrs.contains(idTokenAcrClaim.toString())) {
         this.infoLogNonMatchingAcrClaim(requestedAcrs, idTokenAcrClaim);
         return false;
      } else {
         return true;
      }
   }

   private void infoLogMissingAcrClaim(List<String> requestedClaims) {
      this.getLogger()
         .info(
            "OpenID Connect ID Token validation failed. No acr claim is present but the client requested one of the following: "
               + String.join(", ", requestedClaims)
         );
   }

   private void infoLogNonMatchingAcrClaim(List<String> requestedAcrs, Object acrClaim) {
      this.getLogger()
         .info(
            String.format(
               "OpenID Connect ID Token validation failed. None of the acr claims requested by the client are present. Requested: %s but got: %s",
               String.join(", ", requestedAcrs),
               acrClaim
            )
         );
   }

   private boolean validateIdTokenNotExpired(Map<String, Object> claims) {
      Object exp = claims.get(OpenIdConnectClaim.EXPIRATION_TIME.getClaimName());
      if (!(exp instanceof Date)) {
         this.getLogger()
            .info(
               "OpenID Connect ID Token claim \""
                  + OpenIdConnectClaim.EXPIRATION_TIME.getClaimName()
                  + "\" missing or not a date. The claim value is: \""
                  + exp
                  + "\"."
            );
         return false;
      } else {
         long expTimeStamp = ((Date)exp).getTime();
         if (System.currentTimeMillis() > expTimeStamp) {
            this.getLogger().info("OpenID Connect ID Token is expired. Expiration time stamp is: \"" + expTimeStamp + "\".");
            return false;
         } else {
            this.getLogger().debug("OpenID Connect ID Token contains valid expiration ('exp') claim");
            return true;
         }
      }
   }

   private boolean validateIdTokenIssuer(Map<String, Object> claims) {
      Object iss = claims.get(OpenIdConnectClaim.ISSUER_IDENTIFIER.getClaimName());
      if (iss == null) {
         this.getLogger()
            .info("OpenID Connect ID Token claim \"" + OpenIdConnectClaim.ISSUER_IDENTIFIER.getClaimName() + "\" is invalid. The claim value is: \"null\".");
         return false;
      }

      if (this.settings.customIssuerClaim().isPresent()) {
         if (iss.toString().equals(this.settings.customIssuerClaim().get())) {
            this.getLogger().debug("OpenID Connect ID Token contains valid issuer ('iss') claim");
            return true;
         } else {
            this.getLogger()
               .info(
                  "OpenID Connect ID Token claim \""
                     + OpenIdConnectClaim.ISSUER_IDENTIFIER.getClaimName()
                     + "\" is invalid. The claim value is: \""
                     + iss
                     + "\" and should be \""
                     + this.settings.customIssuerClaim().get()
                     + "\"."
               );
            return false;
         }
      } else {
         try {
            URI issUrl = new URI(iss.toString());
            URI authzServerUrl = this.settings.authorizationEndpointUri().get();
            if (authzServerUrl.getHost().equals(issUrl.getHost())) {
               this.getLogger().debug("OpenID Connect ID Token contains valid issuer ('iss') claim");
               return true;
            } else {
               this.getLogger()
                  .info(
                     "OpenID Connect ID Token claim \""
                        + OpenIdConnectClaim.ISSUER_IDENTIFIER.getClaimName()
                        + "\" is invalid. Reason: Claim value does not match authorization server host. The claim value is: \""
                        + iss
                        + "\"."
                  );
               return false;
            }
         } catch (URISyntaxException e) {
            this.getLogger()
               .info(
                  "OpenID Connect ID Token claim \""
                     + OpenIdConnectClaim.ISSUER_IDENTIFIER.getClaimName()
                     + "\" is invalid. Reason: Claim value could not be parsed as URL. The claim value is: \""
                     + iss
                     + "\".",
                  e
               );
            return false;
         }
      }
   }

   private boolean validateIdTokenAudience(Map<String, Object> claims) {
      if (!this.verifyAudience(claims)) {
         this.getLogger()
            .warn(
               OAuth2LogTypes.OAUTH2_MESSAGE,
               "OpenID Connect ID Token claim \""
                  + OpenIdConnectClaim.AUDIENCE.getClaimName()
                  + "\" does not contain the Client ID "
                  + LogUtil.quote(this.settings.clientId().id())
                  + ". The claim's value is: \""
                  + claims.get(OpenIdConnectClaim.AUDIENCE.getClaimName())
                  + "\"."
            );
         return false;
      } else if (!this.verifyAuthorizingParty(claims)) {
         this.getLogger()
            .warn(
               OAuth2LogTypes.OAUTH2_MESSAGE,
               "OpenID Connect ID Token claim \""
                  + OpenIdConnectClaim.AUTHORIZED_PARTY.getClaimName()
                  + "\" is not set to the Client ID "
                  + LogUtil.quote(this.settings.clientId().id())
                  + ". The claim's value is: \""
                  + claims.get(OpenIdConnectClaim.AUTHORIZED_PARTY.getClaimName())
                  + "\"."
            );
         return false;
      } else {
         this.getLogger().debug("OpenID Connect ID Token contains valid audience ('aud') and authorization party ('azp') claims");
         return true;
      }
   }

   private boolean verifyAudience(Map<String, Object> claims) {
      Object aud = claims.get(OpenIdConnectClaim.AUDIENCE.getClaimName());
      if (aud == null) {
         this.getLogger().info("OpenID Connect ID Token claim \"" + OpenIdConnectClaim.AUDIENCE.getClaimName() + "\" is empty or not present.");
         return false;
      } else {
         return aud instanceof Collection
            ? this.validateIdTokenAudienceCollection((Collection<?>)aud)
            : this.validateSingleIdTokenAudienceOrAzpValue(aud.toString());
      }
   }

   private boolean verifyAuthorizingParty(Map<String, Object> claims) {
      Object azp = claims.get(OpenIdConnectClaim.AUTHORIZED_PARTY.getClaimName());
      return azp == null || this.validateSingleIdTokenAudienceOrAzpValue(azp.toString());
   }

   private boolean validateIdTokenAudienceCollection(Collection<?> aud) {
      for (Object object : aud) {
         if (this.validateSingleIdTokenAudienceOrAzpValue(object.toString())) {
            return true;
         }
      }

      return false;
   }

   private boolean validateSingleIdTokenAudienceOrAzpValue(String aud) {
      if (StringUtils.isBlank(aud)) {
         return false;
      }

      if (AudienceClaimValidationMethod.STANDARD.equals(this.settings.audienceClaimValidationMethod())) {
         return this.settings.clientId().id().equals(aud);
      }

      if (AudienceClaimValidationMethod.CUSTOM.equals(this.settings.audienceClaimValidationMethod())) {
         return this.settings.customAudienceClaim().map(aud::equals).orElse(false);
      }

      if (AudienceClaimValidationMethod.REDIRECT_URL.equals(this.settings.audienceClaimValidationMethod())) {
         try {
            URI audClaimUrl = new URI(aud);
            URI redirectUrl = this.settings.redirectUri();
            if (redirectUrl.getHost().equals(audClaimUrl.getHost())) {
               return true;
            }
         } catch (URISyntaxException e) {
            this.getLogger().info("Unable to parse \"aud\" value " + LogUtil.quote(aud) + " as a URL:" + e.getMessage());
         }
      }

      return false;
   }

   @SuppressFBWarnings(justification = "generated code")
   public static OpenIDConnectClientParserImpl.OpenIDConnectClientParserImplBuilder openIdConnectBuilder() {
      return new OpenIDConnectClientParserImpl.OpenIDConnectClientParserImplBuilder();
   }

   @SuppressFBWarnings(justification = "generated code")
   public static class OpenIDConnectClientParserImplBuilder {
      @SuppressFBWarnings(justification = "generated code")
      private ELogger logger;
      @SuppressFBWarnings(justification = "generated code")
      private Supplier<OpenIDConnectIdTokenSignatureValidator> signatureValidator;
      @SuppressFBWarnings(justification = "generated code")
      private List<ClaimsValidator> additionalClaimValidators;
      @SuppressFBWarnings(justification = "generated code")
      private OpenIDConnectClientParserSettings settings;

      @SuppressFBWarnings(justification = "generated code")
      OpenIDConnectClientParserImplBuilder() {
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIDConnectClientParserImpl.OpenIDConnectClientParserImplBuilder logger(@NonNull ELogger logger) {
         if (logger == null) {
            throw new NullPointerException("logger is marked non-null but is null");
         }

         this.logger = logger;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIDConnectClientParserImpl.OpenIDConnectClientParserImplBuilder signatureValidator(
         @NonNull Supplier<OpenIDConnectIdTokenSignatureValidator> signatureValidator
      ) {
         if (signatureValidator == null) {
            throw new NullPointerException("signatureValidator is marked non-null but is null");
         }

         this.signatureValidator = signatureValidator;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIDConnectClientParserImpl.OpenIDConnectClientParserImplBuilder additionalClaimValidators(
         @NonNull List<ClaimsValidator> additionalClaimValidators
      ) {
         if (additionalClaimValidators == null) {
            throw new NullPointerException("additionalClaimValidators is marked non-null but is null");
         }

         this.additionalClaimValidators = additionalClaimValidators;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIDConnectClientParserImpl.OpenIDConnectClientParserImplBuilder settings(@NonNull OpenIDConnectClientParserSettings settings) {
         if (settings == null) {
            throw new NullPointerException("settings is marked non-null but is null");
         }

         this.settings = settings;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIDConnectClientParserImpl build() {
         return new OpenIDConnectClientParserImpl(this.logger, this.signatureValidator, this.additionalClaimValidators, this.settings);
      }

      @SuppressFBWarnings(justification = "generated code")
      @Override
      public String toString() {
         return "OpenIDConnectClientParserImpl.OpenIDConnectClientParserImplBuilder(logger="
            + this.logger
            + ", signatureValidator="
            + this.signatureValidator
            + ", additionalClaimValidators="
            + this.additionalClaimValidators
            + ", settings="
            + this.settings
            + ")";
      }
   }
}