package com.airlock.iam.common.domain.service.jwt;

import com.airlock.iam.base.misc.json.JsonCoding;
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.encryption.JweToken;
import com.airlock.iam.common.domain.model.jwt.signature.JsonWebSignatureToken;
import com.airlock.iam.common.domain.service.jwt.encryption.JwtTicketEncrypter;
import com.airlock.iam.common.domain.service.jwt.signature.JwtSigner;
import com.airlock.iam.common.domain.service.jwt.signature.JwtSignerException;
import com.airlock.iam.core.misc.util.ticket.KeyMultiValue;
import com.airlock.iam.core.misc.util.ticket.SecureTicketEncoder;
import com.airlock.iam.core.misc.util.ticket.Ticket;
import com.airlock.iam.core.misc.util.ticket.TicketEncodingException;
import com.airlock.iam.elog.api.ELogger;
import com.airlock.iam.elog.api.ELoggerFactory;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.JWTClaimsSet.Builder;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import jakarta.annotation.Nonnull;
import java.util.Arrays;
import java.util.Collections;
import java.util.Date;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.UUID;
import java.util.Map.Entry;
import lombok.NonNull;
import org.apache.commons.collections4.CollectionUtils;
import org.jooq.tools.StringUtils;

public final class JwtTicketEncoderImpl implements SecureTicketEncoder {
   private static final ELogger LOG = ELoggerFactory.getLogger(JwtTicketEncoderImpl.class);
   private static final String SUBJECT_CLAIM = "sub";
   private static final String EXPIRY_CLAIM = "exp";
   private static final String NOT_BEFORE_CLAIM = "nbf";
   private static final String ISSUED_AT_CLAIM = "iat";
   private static final String ISSUER_CLAIM = "iss";
   private static final String AUDIENCE_CLAIM = "aud";
   private static final String JWT_ID_CLAIM = "jti";
   @NonNull
   private final Optional<String> usernameTicketKey;
   @NonNull
   private final JwtSigner signer;
   @NonNull
   private final TimeProvider timeProvider;
   private Integer expirationTime;
   @NonNull
   private Integer validNotBeforeSkew;
   private boolean enforceJwtId;
   @NonNull
   private Optional<String> issuer;
   @NonNull
   private List<String> audience;
   @NonNull
   private List<String> claimsStoredAsArray;
   @NonNull
   private List<String> claimsStoredAsJson;
   @NonNull
   private Optional<JwtTicketEncrypter> encrypter;
   private boolean useTicketExpiryDate;

   private JwtTicketEncoderImpl(@Nonnull JwtSigner signer, @Nonnull TimeProvider timeProvider) {
      this.usernameTicketKey = Optional.empty();
      this.signer = signer;
      this.timeProvider = timeProvider;
      this.validNotBeforeSkew = 0;
      this.enforceJwtId = false;
      this.issuer = Optional.empty();
      this.audience = Collections.emptyList();
      this.claimsStoredAsArray = Collections.emptyList();
      this.claimsStoredAsJson = Collections.emptyList();
      this.encrypter = Optional.empty();
      this.useTicketExpiryDate = true;
   }

   public static JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder builder(String usernameTicketKeyVal, JwtSigner signer, TimeProvider timeProvider) {
      return new JwtTicketEncoderImpl(signer, timeProvider).toBuilder().usernameTicketKey(Optional.of(usernameTicketKeyVal));
   }

   public static JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder builder(JwtSigner signer, TimeProvider timeProvider) {
      return new JwtTicketEncoderImpl(signer, timeProvider).toBuilder();
   }

   @Override
   public String encode(Ticket ticket, String urlEncodingScheme) throws TicketEncodingException {
      Objects.requireNonNull(ticket, "ticket must not be null");
      LOG.debug("Starting to encode the JWT ticket.");
      return this.jwt(ticket);
   }

   private String jwt(Ticket ticket) throws TicketEncodingException {
      JsonWebSignatureToken signedJwt = this.signedJwt(ticket);
      return this.encrypter.isPresent() ? this.encryptedJwt(signedJwt).getJwe() : signedJwt.tokenString();
   }

   private JsonWebSignatureToken signedJwt(Ticket ticket) throws TicketEncodingException {
      try {
         JsonWebSignatureToken signedJwt = this.signer.sign(this.jwtClaimsSetFor(ticket));
         LOG.debug("Successfully signed the JWT ticket.");
         return signedJwt;
      } catch (JwtSignerException e) {
         throw new TicketEncodingException("Could not sign ticket", e);
      }
   }

   private JweToken encryptedJwt(JsonWebSignatureToken signedJwt) throws TicketEncodingException {
      JweToken encryptedJwt = this.encrypter.get().encrypt(signedJwt);
      LOG.debug("Successfully encrypted the JWT ticket.");
      return encryptedJwt;
   }

   private JwtClaimsSet jwtClaimsSetFor(Ticket ticket) throws TicketEncodingException {
      Builder claimSetBuilder = new Builder().notBeforeTime(this.notBeforeTime(ticket)).issueTime(this.issueTime(ticket));
      this.username(ticket).ifPresent(claimSetBuilder::subject);
      this.expirationTime(ticket).ifPresent(claimSetBuilder::expirationTime);
      this.issuer(ticket).ifPresent(claimSetBuilder::issuer);
      this.audience(ticket).ifPresent(claimSetBuilder::audience);
      this.jwtId(ticket).ifPresent(claimSetBuilder::jwtID);
      this.addCustomArrayClaims(ticket, claimSetBuilder);
      this.addCustomJsonClaims(ticket, claimSetBuilder);
      this.addCustomClaims(ticket, claimSetBuilder);
      return JwtClaimsSet.of(claimSetBuilder.build().toString());
   }

   private Optional<String> username(Ticket ticket) throws TicketEncodingException {
      return this.usernameTicketKey.isPresent() ? Optional.of(this.username(ticket, this.usernameTicketKey.get())) : Optional.empty();
   }

   private String username(Ticket ticket, String usernameKey) throws TicketEncodingException {
      verifyTicketValueIsNotMultiValue(ticket, usernameKey);
      verifyTicketValueIsNotMultiValue(ticket, "sub");
      Optional<String> ticketUsername = ticket.getFirstValueForName(usernameKey);
      Optional<String> ticketSubject = ticket.getFirstValueForName("sub");
      if (!"sub".equals(usernameKey) && isConflicting(ticketUsername, ticketSubject)) {
         throw new TicketEncodingException(
            "The value for the subject claim is ambiguous, because the ticket has values for keys '" + usernameKey + "' and 'sub'."
         );
      }

      String username = ticketUsername.orElseThrow(
         () -> new TicketEncodingException("The ticket has no value for the configured username ticket key '" + usernameKey + "'.")
      );
      LOG.debug(String.format("Creating JWT ticket for username key '%s'. The username (subject in the JWT) is '%s'.", usernameKey, username));
      return username;
   }

   private Optional<Date> expirationTime(Ticket ticket) throws TicketEncodingException {
      if (this.useTicketExpiryDate) {
         TicketEncodingException.checkEmpty(
            ticket.getFirstValueForName("exp"),
            "The value for the expiry claim is ambiguous, because the ticket contains a key for the registered claim name 'exp' and also has an expiry date. Please remove this key from the ticket."
         );
         return this.expirationTime != null ? Optional.of(this.configuredExpiryDate()) : Optional.of(ticket.getExpiryDate());
      } else {
         return Optional.empty();
      }
   }

   private Date configuredExpiryDate() {
      return new Date(this.timeProvider.getCurrentTime().getTime() + this.expirationTime.intValue() * 1000L);
   }

   private Date notBeforeTime(Ticket ticket) throws TicketEncodingException {
      TicketEncodingException.checkEmpty(
         ticket.getFirstValueForName("nbf"),
         "The Ticket contains a key for the registered claim 'nbf', however the value for this claim is generated by the plugin"
      );
      return Date.from(this.timeProvider.getCurrentInstant().minusSeconds(this.validNotBeforeSkew.intValue()));
   }

   private Date issueTime(Ticket ticket) throws TicketEncodingException {
      TicketEncodingException.checkEmpty(
         ticket.getFirstValueForName("iat"),
         "The Ticket contains a key for the registered claim 'iat', however the value for this claim is generated by the plugin"
      );
      return this.timeProvider.getCurrentTime();
   }

   private Optional<String> issuer(Ticket ticket) throws TicketEncodingException {
      verifyTicketValueIsNotMultiValue(ticket, "iss");
      Optional<String> ticketIssuer = ticket.getFirstValueForName("iss");
      if (isConflicting(ticketIssuer, this.issuer)) {
         throw new TicketEncodingException(
            "The value for the issuer claim is ambiguous, because it is configured, but also present in the ticket. Please remove the configured issuer or remove the 'iss' key from the ticket."
         );
      } else {
         return ticketIssuer.isPresent() ? ticketIssuer : this.issuer;
      }
   }

   private Optional<List<String>> audience(Ticket ticket) throws TicketEncodingException {
      Optional<List<String>> ticketAudience = Optional.ofNullable(ticket.getValueForName("aud")).map(kmv -> Arrays.asList(kmv.getValues()));
      if (ticketAudience.isPresent() && !this.audience.isEmpty() && !CollectionUtils.isEqualCollection(ticketAudience.get(), this.audience)) {
         throw new TicketEncodingException(
            "The value for the audience claim is ambiguous, because it is configured, but also present in the ticket. Please remove the configured audience or remove the 'aud' key from the ticket."
         );
      } else {
         return ticketAudience.isPresent() ? ticketAudience : optional(this.audience);
      }
   }

   private Optional<String> jwtId(Ticket ticket) throws TicketEncodingException {
      verifyTicketValueIsNotMultiValue(ticket, "jti");
      Optional<String> ticketUniqueId = ticket.getUniqueId();
      Optional<String> ticketJwtId = ticket.getFirstValueForName("jti");
      if (isConflicting(ticketUniqueId, ticketJwtId)) {
         throw new TicketEncodingException("The value for the JWT ID claim is ambiguous, because the ticket defines both a unique id and the key 'jti'.");
      } else {
         return ticketUniqueId.or(() -> ticketJwtId).or(() -> Optional.of(UUID.randomUUID().toString()).filter(s -> this.enforceJwtId));
      }
   }

   private void addCustomArrayClaims(Ticket ticket, Builder claimSetBuilder) {
      for (String claim : this.claimsStoredAsArray) {
         List<String> claimValue = Optional.ofNullable(ticket.getValueForName(claim))
            .map(kmv -> Arrays.asList(kmv.getValues()))
            .orElse(Collections.emptyList());
         LOG.debug("Adding custom claim with key '" + claim + "' to JWT as array.");
         claimSetBuilder.claim(claim, claimValue);
      }
   }

   private void addCustomJsonClaims(Ticket ticket, Builder claimSetBuilder) throws TicketEncodingException {
      for (String claim : this.claimsStoredAsJson) {
         KeyMultiValue kmv = ticket.getValueForName(claim);
         if (Objects.isNull(kmv)) {
            LOG.info(String.format("No value object found for key '%s' - skipping encoding.", claim));
         } else {
            String firstValue = kmv.getFirstValue();
            if (StringUtils.isEmpty(firstValue)) {
               LOG.info(String.format("Empty value found for key '%s' - skipping encoding.", claim));
            } else {
               int nofValues = kmv.getNumberOfValues();
               switch (nofValues) {
                  case 0:
                     LOG.info(String.format("No value found for key '%s' - skipping encoding.", kmv.getKey()));
                     break;
                  case 1:
                     LOG.debug("Adding custom claim with key '" + claim + "' to JWT as JSON.");
                     claimSetBuilder.claim(claim, valueToJson(kmv.getKey(), firstValue));
                     break;
                  default:
                     LOG.info(
                        String.format(
                           "Too many values found for key '%s' (expected 1, received %d) - using first value to encode as JSON.", kmv.getKey(), nofValues
                        )
                     );
                     claimSetBuilder.claim(claim, valueToJson(kmv.getKey(), firstValue));
               }
            }
         }
      }
   }

   private static Object valueToJson(String key, String value) throws TicketEncodingException {
      if ("null".equals(value)) {
         return null;
      } else {
         Optional<Object> json = JsonCoding.tryFromJsonToObject(value);
         if (json.isEmpty()) {
            LOG.debug(String.format("Invalid JSON for key '%s' and value '%s'", key, value));
            throw new TicketEncodingException(String.format("Invalid JSON for key '%s'", key));
         } else {
            return json.get();
         }
      }
   }

   private void addCustomClaims(Ticket ticket, Builder claimSetBuilder) {
      for (Entry<String, KeyMultiValue> entry : ticket.entrySet()) {
         String key = entry.getKey();
         if (!this.isSpecificallyEncodedClaim(key)) {
            KeyMultiValue kmvForName = entry.getValue();
            Object values = kmvForName.getNumberOfValues() == 1 ? kmvForName.getFirstValue() : kmvForName.getValues();
            LOG.debug("Adding custom claim with key '" + key + "' to JWT");
            claimSetBuilder.claim(key, values);
         }
      }
   }

   private boolean isSpecificallyEncodedClaim(String key) {
      return isKeyARegisteredClaim(key) || this.isUsernameKeyOrUniqueIdKey(key) || this.isCustomArrayClaim(key) || this.isCustomJsonClaim(key);
   }

   private static boolean isKeyARegisteredClaim(String ticketKey) {
      return JWTClaimsSet.getRegisteredNames().contains(ticketKey);
   }

   private boolean isUsernameKeyOrUniqueIdKey(String ticketKey) {
      return this.usernameTicketKey.<Boolean>map(t -> Arrays.asList(t, "uniqueId").contains(ticketKey)).orElse("uniqueId".equals(ticketKey));
   }

   private boolean isCustomArrayClaim(String claim) {
      return this.claimsStoredAsArray.contains(claim);
   }

   private boolean isCustomJsonClaim(String claim) {
      return this.claimsStoredAsJson.contains(claim);
   }

   private static void verifyTicketValueIsNotMultiValue(Ticket ticket, String ticketKey) throws TicketEncodingException {
      if (Optional.ofNullable(ticket.getValueForName(ticketKey)).map(KeyMultiValue::getNumberOfValues).orElse(0) > 1) {
         throw new TicketEncodingException("There are multiple values in the ticket for key '" + ticketKey + "' which is not allowed.");
      }
   }

   private static boolean isConflicting(Optional<String> val1, Optional<String> val2) {
      return val1.isPresent() && val2.isPresent() && !val1.equals(val2);
   }

   private static Optional<List<String>> optional(List<String> list) {
      return list.isEmpty() ? Optional.empty() : Optional.of(list);
   }

   @SuppressFBWarnings(justification = "generated code")
   public JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder toBuilder() {
      return new JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder()
         .usernameTicketKey(this.usernameTicketKey)
         .signer(this.signer)
         .timeProvider(this.timeProvider)
         .expirationTime(this.expirationTime)
         .validNotBeforeSkew(this.validNotBeforeSkew)
         .enforceJwtId(this.enforceJwtId)
         .issuer(this.issuer)
         .audience(this.audience)
         .claimsStoredAsArray(this.claimsStoredAsArray)
         .claimsStoredAsJson(this.claimsStoredAsJson)
         .encrypter(this.encrypter)
         .useTicketExpiryDate(this.useTicketExpiryDate);
   }

   @SuppressFBWarnings(justification = "generated code")
   public JwtTicketEncoderImpl(
      @NonNull Optional<String> usernameTicketKey,
      @NonNull JwtSigner signer,
      @NonNull TimeProvider timeProvider,
      Integer expirationTime,
      @NonNull Integer validNotBeforeSkew,
      boolean enforceJwtId,
      @NonNull Optional<String> issuer,
      @NonNull List<String> audience,
      @NonNull List<String> claimsStoredAsArray,
      @NonNull List<String> claimsStoredAsJson,
      @NonNull Optional<JwtTicketEncrypter> encrypter,
      boolean useTicketExpiryDate
   ) {
      if (usernameTicketKey == null) {
         throw new NullPointerException("usernameTicketKey is marked non-null but is null");
      }

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

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

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

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

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

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

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

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

      this.usernameTicketKey = usernameTicketKey;
      this.signer = signer;
      this.timeProvider = timeProvider;
      this.expirationTime = expirationTime;
      this.validNotBeforeSkew = validNotBeforeSkew;
      this.enforceJwtId = enforceJwtId;
      this.issuer = issuer;
      this.audience = audience;
      this.claimsStoredAsArray = claimsStoredAsArray;
      this.claimsStoredAsJson = claimsStoredAsJson;
      this.encrypter = encrypter;
      this.useTicketExpiryDate = useTicketExpiryDate;
   }

   @SuppressFBWarnings(justification = "generated code")
   public static class JwtTicketEncoderImplBuilder {
      @SuppressFBWarnings(justification = "generated code")
      private Optional<String> usernameTicketKey;
      @SuppressFBWarnings(justification = "generated code")
      private JwtSigner signer;
      @SuppressFBWarnings(justification = "generated code")
      private TimeProvider timeProvider;
      @SuppressFBWarnings(justification = "generated code")
      private Integer expirationTime;
      @SuppressFBWarnings(justification = "generated code")
      private Integer validNotBeforeSkew;
      @SuppressFBWarnings(justification = "generated code")
      private boolean enforceJwtId;
      @SuppressFBWarnings(justification = "generated code")
      private Optional<String> issuer;
      @SuppressFBWarnings(justification = "generated code")
      private List<String> audience;
      @SuppressFBWarnings(justification = "generated code")
      private List<String> claimsStoredAsArray;
      @SuppressFBWarnings(justification = "generated code")
      private List<String> claimsStoredAsJson;
      @SuppressFBWarnings(justification = "generated code")
      private Optional<JwtTicketEncrypter> encrypter;
      @SuppressFBWarnings(justification = "generated code")
      private boolean useTicketExpiryDate;

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

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder usernameTicketKey(@NonNull Optional<String> usernameTicketKey) {
         if (usernameTicketKey == null) {
            throw new NullPointerException("usernameTicketKey is marked non-null but is null");
         }

         this.usernameTicketKey = usernameTicketKey;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder signer(@NonNull JwtSigner signer) {
         if (signer == null) {
            throw new NullPointerException("signer is marked non-null but is null");
         }

         this.signer = signer;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder timeProvider(@NonNull TimeProvider timeProvider) {
         if (timeProvider == null) {
            throw new NullPointerException("timeProvider is marked non-null but is null");
         }

         this.timeProvider = timeProvider;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder expirationTime(Integer expirationTime) {
         this.expirationTime = expirationTime;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder validNotBeforeSkew(@NonNull Integer validNotBeforeSkew) {
         if (validNotBeforeSkew == null) {
            throw new NullPointerException("validNotBeforeSkew is marked non-null but is null");
         }

         this.validNotBeforeSkew = validNotBeforeSkew;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder enforceJwtId(boolean enforceJwtId) {
         this.enforceJwtId = enforceJwtId;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder issuer(@NonNull Optional<String> issuer) {
         if (issuer == null) {
            throw new NullPointerException("issuer is marked non-null but is null");
         }

         this.issuer = issuer;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder audience(@NonNull List<String> audience) {
         if (audience == null) {
            throw new NullPointerException("audience is marked non-null but is null");
         }

         this.audience = audience;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder claimsStoredAsArray(@NonNull List<String> claimsStoredAsArray) {
         if (claimsStoredAsArray == null) {
            throw new NullPointerException("claimsStoredAsArray is marked non-null but is null");
         }

         this.claimsStoredAsArray = claimsStoredAsArray;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder claimsStoredAsJson(@NonNull List<String> claimsStoredAsJson) {
         if (claimsStoredAsJson == null) {
            throw new NullPointerException("claimsStoredAsJson is marked non-null but is null");
         }

         this.claimsStoredAsJson = claimsStoredAsJson;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder encrypter(@NonNull Optional<JwtTicketEncrypter> encrypter) {
         if (encrypter == null) {
            throw new NullPointerException("encrypter is marked non-null but is null");
         }

         this.encrypter = encrypter;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder useTicketExpiryDate(boolean useTicketExpiryDate) {
         this.useTicketExpiryDate = useTicketExpiryDate;
         return this;
      }

      @SuppressFBWarnings(justification = "generated code")
      public JwtTicketEncoderImpl build() {
         return new JwtTicketEncoderImpl(
            this.usernameTicketKey,
            this.signer,
            this.timeProvider,
            this.expirationTime,
            this.validNotBeforeSkew,
            this.enforceJwtId,
            this.issuer,
            this.audience,
            this.claimsStoredAsArray,
            this.claimsStoredAsJson,
            this.encrypter,
            this.useTicketExpiryDate
         );
      }

      @SuppressFBWarnings(justification = "generated code")
      @Override
      public String toString() {
         return "JwtTicketEncoderImpl.JwtTicketEncoderImplBuilder(usernameTicketKey="
            + this.usernameTicketKey
            + ", signer="
            + this.signer
            + ", timeProvider="
            + this.timeProvider
            + ", expirationTime="
            + this.expirationTime
            + ", validNotBeforeSkew="
            + this.validNotBeforeSkew
            + ", enforceJwtId="
            + this.enforceJwtId
            + ", issuer="
            + this.issuer
            + ", audience="
            + this.audience
            + ", claimsStoredAsArray="
            + this.claimsStoredAsArray
            + ", claimsStoredAsJson="
            + this.claimsStoredAsJson
            + ", encrypter="
            + this.encrypter
            + ", useTicketExpiryDate="
            + this.useTicketExpiryDate
            + ")";
      }
   }
}