package com.atlassian.asap.api;
   
   import com.atlassian.asap.core.serializer.Json;
   import com.google.common.collect.ImmutableList;
   import com.google.common.collect.ImmutableSet;
   import org.apache.commons.lang3.builder.EqualsBuilder;
   import org.apache.commons.lang3.builder.HashCodeBuilder;
   import org.apache.commons.lang3.builder.ToStringBuilder;
   import org.apache.commons.lang3.builder.ToStringStyle;
  
  import javax.json.JsonArray;
  import javax.json.JsonArrayBuilder;
  import javax.json.JsonNumber;
  import javax.json.JsonObject;
  import javax.json.JsonObjectBuilder;
  import javax.json.JsonString;
  import javax.json.JsonValue;
  import java.io.IOException;
  import java.io.ObjectInputStream;
  import java.io.ObjectOutputStream;
  import java.io.Serializable;
  import java.time.Duration;
  import java.time.Instant;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Optional;
  import java.util.Set;
  import java.util.UUID;
  import java.util.stream.Collectors;
  
  /**
   * A fluent builder for constructing a {@link com.atlassian.asap.api.Jwt} object.
   */
  public final class JwtBuilder {
      /**
       * Token lifetime, unless a specific lifetime is explicitly set. This is the time span between the
       * iat and the exp claims.
       */
      public static final Duration DEFAULT_LIFETIME = Duration.ofSeconds(60);
  
      private SigningAlgorithm alg;
      private String keyId;
      private String iss;
      private Optional<String> sub;
      private List<String> aud;
      private Instant iat;
      private Instant exp;
      private Optional<Instant> nbf;
      private String jti;
      private JsonObject customClaims;
  
      private JwtBuilder() {
          Instant now = Instant.now();
          notBefore(Optional.of(now));
          issuedAt(now);
          expirationTime(now.plus(DEFAULT_LIFETIME));
          jwtId(UUID.randomUUID().toString());
          algorithm(SigningAlgorithm.RS256);
          sub = Optional.empty();
          customClaims = Json.provider().createObjectBuilder().build();
      }
  
      /**
       * Construct a simple jwt builder initialised with default claim values as follows:
       * <ul>
       * <li> nbf, iat claim set to current system time </li>
       * <li> exp claim set current time plus default expiry as defined in {@link JwtBuilder#DEFAULT_LIFETIME} </li>
       * <li> jti claim set to a random UUID. </li>
       * <li> alg header set to {@link SigningAlgorithm#RS256}. </li>
       * </ul>
       *
       * @return a new fluent builder
       */
      public static JwtBuilder newJwt() {
          return new JwtBuilder();
      }
  
      /**
       * Construct a JWT builder initialised from a prototype JWT. Time claims are rebased to the current time,
       * and the jti is set to a random UUID.
       *
       * @param jwtPrototype Jwt to use as prototype
       * @return a new fluent builder
       */
      public static JwtBuilder newFromPrototype(Jwt jwtPrototype) {
          Instant now = Instant.now();
          return copyJwt(jwtPrototype)
                  .notBefore(Optional.of(now))
                  .issuedAt(now)
                  .expirationTime(now.plus(DEFAULT_LIFETIME))
                  .jwtId(UUID.randomUUID().toString());
      }
  
      /**
       * Returns a builder initialised to make an identical copy of the given original Jwt.
       *
       * @param prototype Jwt to use as original
       * @return a Jwt builder initialised to be an identical copy of the original
      */
     public static JwtBuilder copyJwt(Jwt prototype) {
         return new JwtBuilder()
                 .algorithm(prototype.getHeader().getAlgorithm())
                 .keyId(prototype.getHeader().getKeyId())
                 .issuer(prototype.getClaims().getIssuer())
                 .subject(prototype.getClaims().getSubject())
                 .audience(prototype.getClaims().getAudience())
                 .issuedAt(prototype.getClaims().getIssuedAt())
                 .expirationTime(prototype.getClaims().getExpiry())
                 .notBefore(prototype.getClaims().getNotBefore())
                 .jwtId(prototype.getClaims().getJwtId())
                 .customClaims(prototype.getClaims().getJson());
     }
 
     /**
      * Sets the key id jws header.
      *
      * @param keyId the key id for the jws header of this jwt
      * @return the fluent builder
      */
     public JwtBuilder keyId(String keyId) {
         this.keyId = keyId;
         return this;
     }
 
     /**
      * Sets the algorithm (alg) jws header.
      *
      * @param alg the alg for the jws header of this jwt
      * @return the fluent builder
      */
     public JwtBuilder algorithm(SigningAlgorithm alg) {
         this.alg = alg;
         return this;
     }
 
     /**
      * Sets the audience (aud) claim.
      *
      * @param aud an iterable containing one or more audiences for the jwt
      * @return the fluent builder
      */
     public JwtBuilder audience(Iterable<String> aud) {
         this.aud = ImmutableList.copyOf(aud);
         return this;
     }
 
     /**
      * Sets the audience (aud) claim.
      *
      * @param aud one or more audiences for the jwt
      * @return the fluent builder
      */
     public JwtBuilder audience(String... aud) {
         this.aud = ImmutableList.copyOf(aud);
         return this;
     }
 
     /**
      * Sets the expiration time (exp) claim.
      *
      * @param expiry the expiration time
      * @return the fluent builder
      */
     public JwtBuilder expirationTime(Instant expiry) {
         this.exp = expiry;
         return this;
     }
 
     /**
      * Set the issued at (iat) claim.
      *
      * @param iat the issued at time
      * @return the fluent builder
      */
     public JwtBuilder issuedAt(Instant iat) {
         this.iat = iat;
         return this;
     }
 
     /**
      * Set the issuer (iss) claim.
      *
      * @param iss the issuer
      * @return the fluent builder
      */
     public JwtBuilder issuer(String iss) {
         this.iss = iss;
         return this;
     }
 
     /**
      * Set the jwt id (jti) claim.
      *
      * @param jti a unique id for the jwt
      * @return the fluent builder
      */
     public JwtBuilder jwtId(String jti) {
         this.jti = jti;
         return this;
     }
 
     /**
      * Set the not before (nbf) claim.
      *
      * @param nbf the not before date
      * @return the fluent builder
      */
     public JwtBuilder notBefore(Optional<Instant> nbf) {
         this.nbf = nbf;
         return this;
     }
 
     /**
      * Sets the subject (sub) claim for this jwt.
      *
      * @param sub the subject
      * @return the fluent builder
      */
     public JwtBuilder subject(Optional<String> sub) {
         this.sub = sub;
         return this;
     }
 
     /**
      * Sets the custom claims for this jwt. Any previously set custom claim is discarded. If any of the custom
      * claims conflicts with a registered claim, the registered claim takes precedence.
      *
      * @param customClaims the custom claims
      * @return the fluent builder
      */
     public JwtBuilder customClaims(JsonObject customClaims) {
         this.customClaims = customClaims;
         return this;
     }
 
     /**
      * @return a serializable JWT object representing all the values specified to this builder
      * @throws java.lang.NullPointerException if some required parameter has not been specified
      */
     public Jwt build() {
         JwsHeader header = new ImmutableJwsHeader(alg, keyId);
 
         JsonObjectBuilder claimsJsonObjectBuilder = Json.provider().createObjectBuilder();
         // custom claims are added first, so in case in conflict with a registered claim the latter wins
         Set<String> optionalRegisteredClaimKeys = ImmutableSet.of(JwtClaims.RegisteredClaim.SUBJECT.key(), JwtClaims.RegisteredClaim.NOT_BEFORE.key());
         for (Map.Entry<String, JsonValue> entry : customClaims.entrySet()) {
             if (!optionalRegisteredClaimKeys.contains(entry.getKey())) {
                 claimsJsonObjectBuilder.add(entry.getKey(), entry.getValue());
             }
         }
         for (Map.Entry<String, JsonValue> entry : ImmutableJwtClaims.getRegisteredClaims(iss, sub, aud, exp, nbf, iat, jti).entrySet()) {
             claimsJsonObjectBuilder.add(entry.getKey(), entry.getValue());
         }
         JwtClaims claims = new ImmutableJwtClaims(claimsJsonObjectBuilder.build());
 
         return new ImmutableJwt(header, claims);
     }
 
     /**
      * An immutable value object that represents a JWT.
      */
     private static class ImmutableJwt implements Jwt, Serializable {
         private static final long serialVersionUID = 4437693510625284065L;
 
         private final JwsHeader header;
         private final JwtClaims claimsSet;
 
         ImmutableJwt(JwsHeader header, JwtClaims claims) {
             this.header = Objects.requireNonNull(header);
             this.claimsSet = Objects.requireNonNull(claims);
         }
 
         @Override
         public JwsHeader getHeader() {
             return header;
         }
 
         @Override
         public JwtClaims getClaims() {
             return claimsSet;
         }
 
         @Override
         public String toString() {
             return new ToStringBuilder(this, ToStringStyle.SHORT_PREFIX_STYLE)
                     .append("header", header)
                     .append("claims", claimsSet)
                     .toString();
         }
 
         @Override
         public boolean equals(Object o) {
             if (o == null) {
                 return false;
             }
             if (o == this) {
                 return true;
             }
             if (o.getClass() != getClass()) {
                 return false;
             }
             ImmutableJwt rhs = (ImmutableJwt) o;
             return new EqualsBuilder()
                     .append(header, rhs.header)
                     .append(claimsSet, rhs.claimsSet)
                     .isEquals();
         }
 
         @Override
         public int hashCode() {
             return new HashCodeBuilder()
                     .append(header)
                     .append(claimsSet)
                     .hashCode();
         }
     }
 
     /**
      * An immutable value object that represents the information contained in the JWS header.
      */
     private static final class ImmutableJwsHeader implements JwsHeader, Serializable {
         private static final long serialVersionUID = -4791575710345791956L;
 
         private final SigningAlgorithm algorithm;
         private final String keyId;
 
         ImmutableJwsHeader(SigningAlgorithm algorithm, String keyId) {
             this.algorithm = Objects.requireNonNull(algorithm, "JWT header 'alg' cannot be null");
             this.keyId = Objects.requireNonNull(keyId, "JWT header 'kid' cannot be null");
         }
 
         @Override
         public String getKeyId() {
             return keyId;
         }
 
         @Override
         public SigningAlgorithm getAlgorithm() {
             return algorithm;
         }
 
         @Override
         public String toString() {
             return new ToStringBuilder(this, ToStringStyle.SHORT_PREFIX_STYLE)
                     .append(Header.ALGORITHM.key(), algorithm)
                     .append(Header.KEY_ID.key(), keyId)
                     .toString();
         }
 
         @Override
         public boolean equals(Object o) {
             if (o == null) {
                 return false;
             }
             if (o == this) {
                 return true;
             }
             if (o.getClass() != getClass()) {
                 return false;
             }
             ImmutableJwsHeader rhs = (ImmutableJwsHeader) o;
             return new EqualsBuilder()
                     .append(algorithm, rhs.algorithm)
                     .append(keyId, rhs.keyId)
                     .isEquals();
         }
 
         @Override
         public int hashCode() {
             return new HashCodeBuilder()
                     .append(algorithm)
                     .append(keyId)
                     .hashCode();
         }
     }
 
     /**
      * An immutable value object that represents the information contained in the claims (payload) of a JWT.
      */
     private static class ImmutableJwtClaims implements JwtClaims, Serializable {
         private static final long serialVersionUID = 5227306085811054804L;
 
         private transient JsonObject claimsJsonObject;  // cannot be final due to custom serialization
 
         ImmutableJwtClaims(JsonObject claimsJsonObject) {
             this.claimsJsonObject = Objects.requireNonNull(claimsJsonObject);
         }
 
         static JsonObject getRegisteredClaims(String iss,
                                               Optional<String> sub,
                                               List<String> aud,
                                               Instant exp,
                                               Optional<Instant> nbf,
                                               Instant iat,
                                               String jti) {
             Objects.requireNonNull(iss, "JWT claim 'iss' cannot be null");
             Objects.requireNonNull(sub, "JWT claim 'sub' cannot be null (but it can be None)");
             Objects.requireNonNull(aud, "JWT claim 'aud' cannot be null");
             Objects.requireNonNull(iat, "JWT claim 'iat' cannot be null");
             Objects.requireNonNull(exp, "JWT claim 'exp' cannot be null");
             Objects.requireNonNull(nbf, "JWT claim 'nbf' cannot be null (but it can be None)");
             Objects.requireNonNull(jti, "JWT claim 'jit' cannot be null");
 
             JsonObjectBuilder jsonObjectBuilder = Json.provider().createObjectBuilder();
             jsonObjectBuilder.add(RegisteredClaim.ISSUER.key(), iss);
             jsonObjectBuilder.add(RegisteredClaim.ISSUED_AT.key(), iat.getEpochSecond());
             jsonObjectBuilder.add(RegisteredClaim.EXPIRY.key(), exp.getEpochSecond());
             jsonObjectBuilder.add(RegisteredClaim.JWT_ID.key(), jti);
             sub.ifPresent(s -> jsonObjectBuilder.add(RegisteredClaim.SUBJECT.key(), s));
             nbf.ifPresent(n -> jsonObjectBuilder.add(RegisteredClaim.NOT_BEFORE.key(), n.getEpochSecond()));
             if (aud.size() == 1) { // optimise JSON in case of a single audience
                 jsonObjectBuilder.add(RegisteredClaim.AUDIENCE.key(), aud.iterator().next());
             } else {
                 JsonArrayBuilder audienceArrayBuilder = Json.provider().createArrayBuilder();
                 aud.forEach(audienceArrayBuilder::add);
                 jsonObjectBuilder.add(RegisteredClaim.AUDIENCE.key(), audienceArrayBuilder);
             }
             return jsonObjectBuilder.build();
         }
 
         @Override
         public String getIssuer() {
             return claimsJsonObject.getString(RegisteredClaim.ISSUER.key());
         }
 
         @Override
         public Optional<String> getSubject() {
             return Optional.ofNullable(claimsJsonObject.getString(RegisteredClaim.SUBJECT.key(), null));
         }
 
         @Override
         public Set<String> getAudience() {
             JsonValue jsonValue = claimsJsonObject.getOrDefault(RegisteredClaim.AUDIENCE.key(), Json.provider().createArrayBuilder().build());
             switch (jsonValue.getValueType()) {
                 case ARRAY:
                     return ((JsonArray) jsonValue).getValuesAs(JsonString.class).stream().map(JsonString::getString).collect(Collectors.toSet());
                 case STRING:
                     return ImmutableSet.of(((JsonString) jsonValue).getString());
                 default:
                     throw new RuntimeException("Unexpected 'aud' claim type");
             }
         }
 
         @Override
         public Instant getExpiry() {
             return Instant.ofEpochSecond(claimsJsonObject.getJsonNumber(RegisteredClaim.EXPIRY.key()).longValueExact());
         }
 
         @Override
         public Optional<Instant> getNotBefore() {
             return Optional.ofNullable(claimsJsonObject.getJsonNumber(RegisteredClaim.NOT_BEFORE.key())).map(JsonNumber::longValueExact).map(Instant::ofEpochSecond);
         }
 
         @Override
         public Instant getIssuedAt() {
             return Instant.ofEpochSecond(claimsJsonObject.getJsonNumber(RegisteredClaim.ISSUED_AT.key()).longValueExact());
         }
 
         @Override
         public String getJwtId() {
             return claimsJsonObject.getString(RegisteredClaim.JWT_ID.key());
         }
 
         @Override
         public JsonObject getJson() {
             return claimsJsonObject;
         }
 
         @Override
         public String toString() {
             return new ToStringBuilder(this, ToStringStyle.SHORT_PREFIX_STYLE)
                     .append(RegisteredClaim.ISSUER.key(), getIssuer())
                     .append(RegisteredClaim.SUBJECT.key(), getSubject())
                     .append(RegisteredClaim.AUDIENCE.key(), getAudience())
                     .append(RegisteredClaim.ISSUED_AT.key(), getIssuedAt())
                     .append(RegisteredClaim.EXPIRY.key(), getExpiry())
                     .append(RegisteredClaim.NOT_BEFORE.key(), getNotBefore())
                     .append(RegisteredClaim.JWT_ID.key(), getJwtId())
                     .toString();
         }
 
         @Override
         public boolean equals(Object o) {
             if (o == null) {
                 return false;
             }
             if (o == this) {
                 return true;
             }
             if (o.getClass() != getClass()) {
                 return false;
             }
             ImmutableJwtClaims rhs = (ImmutableJwtClaims) o;
             return claimsJsonObject.equals(rhs.claimsJsonObject);
         }
 
         @Override
         public int hashCode() {
             return claimsJsonObject.hashCode();
         }
 
         private void readObject(ObjectInputStream inputStream) throws ClassNotFoundException, IOException {
             inputStream.defaultReadObject();
             this.claimsJsonObject = Json.provider().createReader(inputStream).readObject();
         }
 
         private void writeObject(ObjectOutputStream outputStream) throws IOException {
             outputStream.defaultWriteObject();
             Json.provider().createWriter(outputStream).writeObject(this.claimsJsonObject);
         }
     }
 }