package com.airlock.iam.authentication.domain.service.rememberme;

import com.airlock.iam.authentication.domain.model.rememberme.RememberMeJspSecret;
import com.airlock.iam.authentication.domain.model.rememberme.RememberMeJspToken;
import com.airlock.iam.authentication.domain.model.rememberme.RememberMeJspTokenHmacKey;
import com.airlock.iam.authentication.domain.model.rememberme.RememberMeJspTokenHmacKeyId;
import com.airlock.iam.base.api.domain.model.user.UserId;
import com.airlock.iam.base.misc.log.LogUtil;
import com.airlock.iam.base.misc.time.RealTimeProvider;
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.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.airlock.iam.flow.shared.domain.model.rememberme.RememberMeTokenLifetime;
import java.util.Optional;
import org.apache.commons.lang3.StringUtils;

public final class RememberMeJspTokenValidator {
   public static final String JWT_CLAIM_INITIALLY_GENERATED_AT = "iga";
   public static final JwsAlgorithmHmac DEFAULT_ALGORITHM = JwsAlgorithmHmac.HS512;
   private static final ELogger LOG = ELoggerFactory.getLogger(RememberMeJspTokenValidator.class);

   private RememberMeJspTokenValidator() {
      throw new AssertionError("utility class");
   }

   public static Optional<RememberMeJspToken> validateToken(
      RememberMeJspSecret secret, RememberMeJspTokenHmacKeyId keyId, RememberMeJspTokenHmacKey hmacKey, RememberMeTokenLifetime lifetime
   ) {
      JsonWebSignatureToken jwsToken = parseSignedJwt(secret.getRawJwt());
      String algorithmFromJWT = jwsToken.header().algorithmName();
      if (!algorithmFromJWT.equals(JwsAlgorithmHmac.HS512.name())) {
         LOG.info(
            String.format("Unexpected hash signature algorithm: %s ; expected %s", LogUtil.quote(algorithmFromJWT), LogUtil.quote(DEFAULT_ALGORITHM.name()))
         );
         return Optional.empty();
      } else {
         JwsVerifier verifier = JwsVerifierImpl.hmacJwsVerifier(hmacKey.getKey().bytes());
         if (!verifier.verify(jwsToken)) {
            LOG.info("Signature of JSP Remember-Me Cookie could not be verified.");
            return Optional.empty();
         } else {
            Optional<RememberMeJspToken> token = mapToJspToken(jwsToken, keyId, jwsToken.header().keyId());
            return token.filter(t -> !t.isExpired(lifetime));
         }
      }
   }

   public static Optional<RememberMeJspTokenHmacKeyId> extractHmacKeyId(RememberMeJspSecret secret) {
      try {
         JsonWebSignatureToken signedJwt = parseSignedJwt(secret.getRawJwt());
         return signedJwt.header().keyId().map(KeyId::keyId).map(RememberMeJspTokenHmacKeyId::rememberMeHmacKeyId);
      } catch (IllegalArgumentException e) {
         LOG.info(String.format("Failed to extract keyId: %s", e.getMessage()));
         return Optional.empty();
      }
   }

   private static JsonWebSignatureToken parseSignedJwt(String jwt) throws IllegalArgumentException {
      JsonWebSignatureToken jwsToken = JsonWebSignatureToken.parse(jwt)
         .orElseThrow(() -> new IllegalArgumentException(String.format("Bad or unsigned JWT %s", LogUtil.logSafeQuote(jwt))));
      JsonWebSignatureToken.Header header = jwsToken.header();
      if (header.type().map(type -> !JsonWebSignatureToken.Header.Type.JWT.equals(type)).orElse(true)) {
         throw new IllegalArgumentException(String.format("No \"typ\" header in JWT: %s", LogUtil.logSafeQuote(jwt)));
      } else {
         return jwsToken;
      }
   }

   private static Optional<RememberMeJspToken> mapToJspToken(JsonWebSignatureToken signedJwt, RememberMeJspTokenHmacKeyId verifiedUserId, Optional<KeyId> keyId) {
      String user = verifiedUserId.getId();
      JwtClaimsSet claims = signedJwt.claimsSet();
      if (keyId.map(KeyId::keyId).<Boolean>map(StringUtils::isBlank).orElse(true)) {
         LOG.debug("RememberMeJspTokenValidator: No keyId in JWT Header, ignoring.");
         return Optional.empty();
      } else if (!user.equals(keyId.orElseThrow().keyId())) {
         LOG.debug(String.format("RememberMeJspTokenValidator: keyId header in JWT doesn't match expected username %s, ignoring.", LogUtil.quote(user)));
         return Optional.empty();
      } else if (claims.subject().<Boolean>map(StringUtils::isBlank).orElse(true)) {
         LOG.debug("RememberMeJspTokenValidator: No subject (sub) in JWT, ignoring.");
         return Optional.empty();
      } else if (!user.equals(claims.subject().orElseThrow())) {
         LOG.debug(String.format("RememberMeJspTokenValidator: Subject (sub) in JWT doesn't match expected username %s, ignoring.", LogUtil.quote(user)));
         return Optional.empty();
      } else if (claims.claimAsInstant("iga").isEmpty()) {
         LOG.debug(String.format("RememberMeJspTokenValidator: No claim %s in JWT, ignoring.", "iga"));
         return Optional.empty();
      } else if (claims.issueTime().isEmpty()) {
         LOG.debug("RememberMeJspTokenValidator: No issue time (iat) in JWT, ignoring.");
         return Optional.empty();
      } else {
         return Optional.of(
            RememberMeJspToken.builder()
               .iamUserId(UserId.userId(verifiedUserId.getId()))
               .generationTimestamp(claims.claimAsInstant("iga"))
               .lastUsedTimestamp(claims.issueTime())
               .time(new RealTimeProvider())
               .build()
         );
      }
   }
}