package com.airlock.iam.oauth2.domain.model.token;

import com.airlock.iam.base.api.domain.model.exception.ViolationException;
import com.airlock.iam.common.domain.model.jwt.signature.JsonWebSignatureToken;
import com.airlock.iam.common.domain.model.jwt.signature.JwsAlgorithm;
import com.airlock.iam.common.domain.model.jwt.signature.JwsAlgorithmHmac;
import com.airlock.iam.common.domain.model.jwt.signature.JwsVerifier;
import com.airlock.iam.common.domain.model.jwt.signature.JwsVerifierImpl;
import com.airlock.iam.common.domain.model.jwt.signature.KeyId;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jwt.JWT;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.JWTParser;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.nio.charset.StandardCharsets;
import java.text.ParseException;
import java.util.Map;
import java.util.Optional;

public final class IdToken {
   private final ELogger logger;
   private final String originalJwtString;
   private final JWT jwt;
   private final JWTClaimsSet claimsSet;
   private final Optional<JsonWebSignatureToken> idToken;

   public static IdToken of(String tokenString) {
      return of(tokenString, ELoggerFactory.getLogger(IdToken.class));
   }

   public static IdToken of(String tokenString, ELogger logger) {
      try {
         JWT jwt = JWTParser.parse(tokenString);
         return new IdToken(logger, tokenString, jwt, jwt.getJWTClaimsSet(), JsonWebSignatureToken.parse(tokenString));
      } catch (ParseException e) {
         throw new ViolationException("Unable to parse JWT", e);
      }
   }

   public static Optional<IdToken> parse(String tokenString, ELogger logger) {
      try {
         return Optional.of(of(tokenString, logger));
      } catch (ViolationException e) {
         return Optional.empty();
      }
   }

   public Optional<String> keyID() {
      return this.idToken.<KeyId>flatMap(jwsToken -> jwsToken.header().keyId()).map(KeyId::keyId);
   }

   public Map<String, Object> claims() {
      return this.claimsSet.getClaims();
   }

   public String tokenString() {
      return this.originalJwtString;
   }

   public boolean validateHS256Signature(String key) {
      if (!this.isSigned()) {
         this.logger.debug("OpenID Connect ID Token HS256 signature validation unsuccessful; ID Token not signed.");
         return false;
      }

      if (!this.validateExpectedAlgorithm(JwsAlgorithmHmac.HS256)) {
         return false;
      }

      try {
         return this.validateSignature(JwsVerifierImpl.hmacJwsVerifier(key.getBytes(StandardCharsets.UTF_8)), JWSAlgorithm.HS256.getName());
      } catch (IllegalArgumentException e) {
         this.logger.info("Invalid key for validating HS256 signature: " + e.getMessage());
         return false;
      }
   }

   public boolean validateMACSignature(String key) {
      try {
         JwsVerifier verifier = JwsVerifierImpl.hmacJwsVerifier(key.getBytes(StandardCharsets.UTF_8));
         if (verifier.supportedJWSAlgorithms().stream().noneMatch(jwsAlgorithm -> jwsAlgorithm.name().equals(this.signatureAlgorithmName()))) {
            this.logger.info("OpenID Connect ID Token MAC signature validation unsuccessful; unsupported algorithm.");
            return false;
         } else {
            return this.validateSignature(verifier, this.signatureAlgorithmName());
         }
      } catch (IllegalArgumentException e) {
         this.logger.info("Invalid key for validating MAC signature: " + e.getMessage());
         return false;
      }
   }

   private boolean validateExpectedAlgorithm(JwsAlgorithm algorithm) {
      boolean algorithmMatches = algorithm.name().equals(this.signatureAlgorithmName());
      if (!algorithmMatches) {
         this.logger.debug("OpenID Connect ID Token signed using " + this.signatureAlgorithmName() + " instead of " + algorithm + ".");
      }

      return algorithmMatches;
   }

   private boolean validateSignature(JwsVerifier verifier, String algorithmName) {
      try {
         if (verifier.verify(this.idToken.orElseThrow(() -> new ViolationException("ID Token signature not present.")))) {
            this.logger.debug("Successfully validated OpenID Connect ID Token signature.");
            return true;
         } else {
            this.logger.debug("OpenID Connect ID Token " + algorithmName + " signature validation unsuccessful.");
            return false;
         }
      } catch (ViolationException e) {
         this.logger.info("Invalid ID Token, could not verify " + algorithmName + " signature: " + e.getMessage());
         return false;
      }
   }

   public String signatureAlgorithmName() {
      return this.jwt.getHeader().getAlgorithm().getName();
   }

   public boolean isSigned() {
      return this.idToken.isPresent();
   }

   public Optional<JsonWebSignatureToken> getSignedJWT() {
      return this.idToken;
   }

   @SuppressFBWarnings(justification = "generated code")
   private IdToken(ELogger logger, String originalJwtString, JWT jwt, JWTClaimsSet claimsSet, Optional<JsonWebSignatureToken> idToken) {
      this.logger = logger;
      this.originalJwtString = originalJwtString;
      this.jwt = jwt;
      this.claimsSet = claimsSet;
      this.idToken = idToken;
   }

   @SuppressFBWarnings(justification = "generated code")
   @Override
   public boolean equals(Object o) {
      if (o == this) {
         return true;
      } else if (!(o instanceof IdToken other)) {
         return false;
      } else {
         Object this$originalJwtString = this.originalJwtString;
         Object other$originalJwtString = other.originalJwtString;
         return this$originalJwtString == null ? other$originalJwtString == null : this$originalJwtString.equals(other$originalJwtString);
      }
   }

   @SuppressFBWarnings(justification = "generated code")
   @Override
   public int hashCode() {
      int PRIME = 59;
      int result = 1;
      Object $originalJwtString = this.originalJwtString;
      return result * 59 + ($originalJwtString == null ? 43 : $originalJwtString.hashCode());
   }
}