package com.airlock.iam.common.domain.model.jwk;

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.jwk.JWK;
import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.KeyUse;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.text.ParseException;
import java.util.List;
import java.util.Optional;

public final class JsonWebKeySet {
   private static final ELogger LOG = ELoggerFactory.getLogger(JsonWebKeySet.class);
   private final JWKSet jwkSet;

   public static Optional<JsonWebKeySet> fromJson(String json) {
      try {
         return Optional.of(new JsonWebKeySet(JWKSet.parse(json)));
      } catch (ParseException e) {
         LOG.info("Failed to parse JWKS from json");
         return Optional.empty();
      }
   }

   public Optional<JsonWebKey> getMatchingKey(Optional<KeyId> keyId) {
      Optional<JsonWebKey> jwk = keyId.map(this::getKeyByMatchingID).orElseGet(this::getSingleSignatureJWK);
      if (jwk.isEmpty()) {
         LOG.info("No matching signature key was found for JWS token.");
      }

      return jwk;
   }

   public String toJson() {
      return this.jwkSet.toString(false);
   }

   private Optional<JsonWebKey> getSingleSignatureJWK() {
      List<JWK> keys = this.jwkSet.getKeys();
      if (keys.size() != 1) {
         LOG.info(String.format("Unable to get single key from JWKS: expected to find precisely 1 JWK but found %d instead.", keys.size()));
         return Optional.empty();
      }

      JWK jwk = keys.getFirst();
      if (KeyUse.SIGNATURE.equals(jwk.getKeyUse())) {
         return fromJWK(jwk);
      }

      LOG.info("JWKS contains a single JWK but its use is not signing.");
      return Optional.empty();
   }

   private Optional<JsonWebKey> getKeyByMatchingID(KeyId keyId) {
      Optional<JWK> key = Optional.ofNullable(this.jwkSet.getKeyByKeyId(keyId.keyId()));
      if (key.isEmpty()) {
         LOG.info(String.format("No matching JWK was found for key ID %s.", keyId));
         return Optional.empty();
      } else {
         JWK jwk = key.get();
         return fromJWK(jwk);
      }
   }

   private static Optional<JsonWebKey> fromJWK(JWK jwk) {
      return switch (jwk.getKeyType().getValue()) {
         case "EC" -> Optional.of(ECJsonWebKey.fromJSON(jwk.toJSONString()));
         case "RSA" -> Optional.of(RSAJsonWebKey.fromJSON(jwk.toJSONString()));
         default -> {
            LOG.info("Signing key type is not supported: Must be asymmetric (RSA/EC) key");
            yield Optional.empty();
         }
      };
   }

   @SuppressFBWarnings(justification = "generated code")
   private JsonWebKeySet(JWKSet jwkSet) {
      this.jwkSet = jwkSet;
   }
}