package com.airlock.iam.oauth2.infrastructure.service.clientauthentication;

import com.airlock.iam.base.misc.log.LogUtil;
import com.airlock.iam.base.misc.time.TimeProvider;
import com.airlock.iam.common.domain.model.jwt.JwtClaimsSet;
import com.airlock.iam.common.domain.model.jwt.signature.JsonWebSignatureToken;
import com.airlock.iam.common.domain.model.jwt.signature.JwsVerifier;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.airlock.iam.oauth2.api.domain.model.OAuth2ClientId;
import com.airlock.iam.oauth2.api.domain.model.clientauthentication.OAuth2AuthClientId;
import com.airlock.iam.oauth2.api.domain.model.clientauthentication.OAuth2ClientAuthentication;
import com.airlock.iam.oauth2.api.domain.service.clientauthentication.OAuth2ClientAuthenticationVerifier;
import com.airlock.iam.oauth2.domain.model.client.OAuth2Client;
import com.airlock.iam.oauth2.domain.model.client.OAuth2ClientLoader;
import com.airlock.iam.oauth2.domain.model.clientauthentication.OAuth2AcceptedClientAssertion;
import com.airlock.iam.oauth2.domain.model.clientauthentication.OAuth2AcceptedClientAssertionId;
import com.airlock.iam.oauth2.domain.model.clientauthentication.OAuth2AcceptedClientAssertionRepository;
import com.airlock.iam.oauth2.domain.model.clientauthentication.OpenIdConnectPrivateKeyJwtAuthenticationPayload;
import com.airlock.iam.oauth2.domain.model.metadata.OAuth2EndpointUriComposer;
import com.nimbusds.jose.JOSEException;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.net.URI;
import java.text.ParseException;
import java.time.Instant;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.stream.Stream;
import org.apache.commons.lang3.Strings;

public class OpenIdConnectPrivateKeyJwtAuthenticationVerifier implements OAuth2ClientAuthenticationVerifier<OpenIdConnectPrivateKeyJwtAuthenticationPayload> {
   private static final ELogger LOG = ELoggerFactory.getLogger(OpenIdConnectPrivateKeyJwtAuthenticationVerifier.class);
   private static final String PRIVATE_KEY_JWT_ASSERTION_TYPE = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer";
   private final OAuth2ClientLoader clientLoader;
   private final TimeProvider timeProvider;
   private final OAuth2EndpointUriComposer uriComposer;
   private final OAuth2AcceptedClientAssertionRepository jtiRepository;

   public Optional<OAuth2ClientAuthentication> verify(OpenIdConnectPrivateKeyJwtAuthenticationPayload authenticationPayload) {
      try {
         if (!"urn:ietf:params:oauth:client-assertion-type:jwt-bearer".equals(authenticationPayload.assertionType())) {
            LOG.debug("Cannot accept client assertion type: " + LogUtil.logSafeQuote(authenticationPayload.assertionType()));
            return Optional.empty();
         }

         JsonWebSignatureToken jwt = new JsonWebSignatureToken(authenticationPayload.assertion());
         OAuth2AuthClientId clientId = authenticationPayload.clientId().orElseThrow();
         if (authenticationPayload.clientIdFromRequest().isPresent() && !Objects.equals(clientId.id(), authenticationPayload.clientIdFromRequest().get())) {
            throw new OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException(
               "The subject (iss/sub) claim does not match the client_id parameter."
            );
         }

         this.verifySignature(jwt, clientId);
         this.verifyClaims(jwt.claimsSet());
         return Optional.of(OAuth2ClientAuthentication.clientAuthentication(authenticationPayload));
      } catch (ParseException e) {
         LOG.info("Could not parse provided JWT. Reason: " + e.getMessage());
      } catch (OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException e) {
         LOG.info("Could not accept token for authentication. Reason: " + e.getMessage());
      } catch (JOSEException | IllegalStateException e) {
         LOG.info("The provided token could not be verified due to an invalid signature.");
      }

      return Optional.empty();
   }

   private void verifyClaims(JwtClaimsSet claims) throws OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException {
      String subject = claims.subject()
         .orElseThrow(() -> new OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException("The required sub (subject) claim is missing."));
      String issuer = claims.issuer()
         .orElseThrow(() -> new OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException("The required iss (issuer) claim is missing."));
      List<String> audience = claims.audience()
         .orElseThrow(() -> new OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException("The required aud (audience) claim is missing."));
      String jwtId = claims.jwtId()
         .orElseThrow(() -> new OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException("The required jti (JWT ID) claim is missing."));
      Instant expires = claims.expirationTime()
         .orElseThrow(() -> new OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException("The required exp (expiration time) claim is missing."));
      if (expires.isBefore(this.timeProvider.getCurrentInstant())) {
         throw new OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException("The token has expired.");
      }

      if (claims.notBeforeTime().isPresent()) {
         boolean nbfValid = claims.notBeforeTime().get().isBefore(this.timeProvider.getCurrentInstant());
         if (!nbfValid) {
            throw new OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException(
               "The not before (nbf) claim does not allow processing of the token at this time."
            );
         }
      }

      if (!Objects.equals(subject, issuer)) {
         throw new OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException("The issuer (iss) and subject (sub) claims are not identical.");
      }

      if (!this.validAudience(audience)) {
         throw new OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException("The audience (aud) claim does not match any valid endpoints.");
      }

      if (!this.isNew(jwtId)) {
         throw new OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException("The JWT ID (jti) claim is not unique.");
      }

      this.persistJti(jwtId, expires);
   }

   private void persistJti(String jwtId, Instant expires) {
      this.jtiRepository.insert(new OAuth2AcceptedClientAssertion(OAuth2AcceptedClientAssertionId.clientAssertionId(jwtId), expires));
   }

   private boolean validAudience(List<String> audience) {
      List<String> validAudiences = Stream.of(
            this.uriComposer.tokenEndpointURI(), this.uriComposer.pushedAuthorizationRequestURI(), this.uriComposer.providerBaseURI()
         )
         .flatMap(OpenIdConnectPrivateKeyJwtAuthenticationVerifier::withAndWithoutTrailingSlash)
         .toList();
      return audience.stream().anyMatch(validAudiences::contains);
   }

   private static Stream<String> withAndWithoutTrailingSlash(URI uri) {
      String uriString = Strings.CS.removeEnd(uri.toASCIIString(), "/");
      return Stream.of(uriString, uriString + "/");
   }

   private boolean isNew(String jwtId) {
      return !this.jtiRepository.contains(OAuth2AcceptedClientAssertionId.clientAssertionId(jwtId));
   }

   private void verifySignature(JsonWebSignatureToken jwt, OAuth2AuthClientId clientId) throws JOSEException, OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException {
      if (!this.verifierFor(clientId, jwt).verify(jwt)) {
         throw new OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException("The JWT signature is invalid.");
      }
   }

   private JwsVerifier verifierFor(OAuth2AuthClientId clientId, JsonWebSignatureToken jwt) throws OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException {
      return this.loadClient(clientId)
         .verifierFor(jwt)
         .orElseThrow(
            () -> new OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException(
               "The client with client ID " + LogUtil.logSafeQuote(clientId.id()) + " cannot verify the given JWS."
            )
         );
   }

   private OAuth2Client loadClient(OAuth2AuthClientId clientId) throws OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException {
      return this.clientLoader
         .tryLoadClient(OAuth2ClientId.oAuth2ClientId(clientId.id()))
         .orElseThrow(
            () -> new OpenIdConnectPrivateKeyJwtAuthenticationVerifier.VerificationException(
               "The client with client ID " + LogUtil.logSafeQuote(clientId.id()) + " was not found."
            )
         );
   }

   @SuppressFBWarnings(justification = "generated code")
   public static OpenIdConnectPrivateKeyJwtAuthenticationVerifier.OpenIdConnectPrivateKeyJwtAuthenticationVerifierBuilder builder() {
      return new OpenIdConnectPrivateKeyJwtAuthenticationVerifier.OpenIdConnectPrivateKeyJwtAuthenticationVerifierBuilder();
   }

   @SuppressFBWarnings(justification = "generated code")
   private OpenIdConnectPrivateKeyJwtAuthenticationVerifier(
      OAuth2ClientLoader clientLoader, TimeProvider timeProvider, OAuth2EndpointUriComposer uriComposer, OAuth2AcceptedClientAssertionRepository jtiRepository
   ) {
      this.clientLoader = clientLoader;
      this.timeProvider = timeProvider;
      this.uriComposer = uriComposer;
      this.jtiRepository = jtiRepository;
   }

   @SuppressFBWarnings(justification = "generated code")
   public static class OpenIdConnectPrivateKeyJwtAuthenticationVerifierBuilder {
      @SuppressFBWarnings(justification = "generated code")
      private OAuth2ClientLoader clientLoader;
      @SuppressFBWarnings(justification = "generated code")
      private TimeProvider timeProvider;
      @SuppressFBWarnings(justification = "generated code")
      private OAuth2EndpointUriComposer uriComposer;
      @SuppressFBWarnings(justification = "generated code")
      private OAuth2AcceptedClientAssertionRepository jtiRepository;

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

      @SuppressFBWarnings(justification = "generated code")
      public OpenIdConnectPrivateKeyJwtAuthenticationVerifier.OpenIdConnectPrivateKeyJwtAuthenticationVerifierBuilder clientLoader(
         OAuth2ClientLoader clientLoader
      ) {
         this.clientLoader = clientLoader;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIdConnectPrivateKeyJwtAuthenticationVerifier.OpenIdConnectPrivateKeyJwtAuthenticationVerifierBuilder timeProvider(TimeProvider timeProvider) {
         this.timeProvider = timeProvider;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIdConnectPrivateKeyJwtAuthenticationVerifier.OpenIdConnectPrivateKeyJwtAuthenticationVerifierBuilder uriComposer(
         OAuth2EndpointUriComposer uriComposer
      ) {
         this.uriComposer = uriComposer;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIdConnectPrivateKeyJwtAuthenticationVerifier.OpenIdConnectPrivateKeyJwtAuthenticationVerifierBuilder jtiRepository(
         OAuth2AcceptedClientAssertionRepository jtiRepository
      ) {
         this.jtiRepository = jtiRepository;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public OpenIdConnectPrivateKeyJwtAuthenticationVerifier build() {
         return new OpenIdConnectPrivateKeyJwtAuthenticationVerifier(this.clientLoader, this.timeProvider, this.uriComposer, this.jtiRepository);
      }

      @SuppressFBWarnings(justification = "generated code")
      @Override
      public String toString() {
         return "OpenIdConnectPrivateKeyJwtAuthenticationVerifier.OpenIdConnectPrivateKeyJwtAuthenticationVerifierBuilder(clientLoader="
            + this.clientLoader
            + ", timeProvider="
            + this.timeProvider
            + ", uriComposer="
            + this.uriComposer
            + ", jtiRepository="
            + this.jtiRepository
            + ")";
      }
   }

   private static class VerificationException extends Exception {
      private VerificationException(String message) {
         super(message);
      }
   }
}