package com.airlock.iam.common.domain.model.jwt.signature;

import com.airlock.iam.base.api.domain.model.exception.ViolationException;
import com.airlock.iam.common.domain.model.jwt.JwtClaimsSet;
import com.nimbusds.jose.JOSEObjectType;
import com.nimbusds.jose.JWSHeader;
import com.nimbusds.jwt.SignedJWT;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.text.ParseException;
import java.util.Optional;
import lombok.NonNull;

public class JsonWebSignatureToken {
   @NonNull
   private final SignedJWT signedJWT;

   public JsonWebSignatureToken(String tokenString) throws ParseException {
      this.signedJWT = SignedJWT.parse(tokenString);
   }

   public static Optional<JsonWebSignatureToken> parse(String tokenString) {
      try {
         return Optional.of(new JsonWebSignatureToken(tokenString));
      } catch (ParseException e) {
         return Optional.empty();
      }
   }

   public static JsonWebSignatureToken of(String tokenString) {
      try {
         return new JsonWebSignatureToken(tokenString);
      } catch (ParseException e) {
         throw new ViolationException("Failed to create signed JWT", e);
      }
   }

   public String tokenString() {
      return this.signedJWT.serialize();
   }

   public JsonWebSignatureToken.Header header() {
      JWSHeader header = this.signedJWT.getHeader();
      return new JsonWebSignatureToken.Header(
         header.getAlgorithm().getName(),
         Optional.ofNullable(header.getKeyID()).map(KeyId::new),
         Optional.ofNullable(header.getType()).<String>map(JOSEObjectType::getType).map(JsonWebSignatureToken.Header.Type::new)
      );
   }

   public JwtClaimsSet claimsSet() {
      return Optional.ofNullable(this.signedJWT.getPayload().toJSONObject())
         .map(JwtClaimsSet::fromMap)
         .orElseThrow(() -> new ViolationException("JWT payload is missing or could not be parsed"));
   }

   @SuppressFBWarnings(justification = "generated code")
   private JsonWebSignatureToken(@NonNull SignedJWT signedJWT) {
      if (signedJWT == null) {
         throw new NullPointerException("signedJWT is marked non-null but is null");
      }

      this.signedJWT = signedJWT;
   }

   public record Header(String algorithmName, Optional<KeyId> keyId, Optional<JsonWebSignatureToken.Header.Type> type) {
      public JwsAlgorithm getValidatedAlgorithm() {
         return JwsAlgorithm.JWS_ALGORITHMS
            .stream()
            .filter(alg -> alg.name().equals(this.algorithmName))
            .findFirst()
            .orElseThrow(() -> new ViolationException("Unsupported algorithm found: " + this.algorithmName));
      }

      public record Type(String type) {
         public static final JsonWebSignatureToken.Header.Type JWT = new JsonWebSignatureToken.Header.Type("JWT");
      }
   }
}