package riomhaire.lti.core.adapters.token;

import com.auth0.jwt.JWT;
import org.jose4j.jwk.HttpsJwks;
import org.jose4j.jwt.JwtClaims;
import org.jose4j.jwt.consumer.InvalidJwtException;
import org.jose4j.jwt.consumer.JwtConsumerBuilder;
import org.jose4j.keys.resolvers.HttpsJwksVerificationKeyResolver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import riomhaire.lti.core.model.interfaces.DecodeException;
import riomhaire.lti.core.model.interfaces.Decoder;

import java.util.Map;

/**
 * This class is a helper class which takes a JWT token an verifies it is valid based on the jwk link provided
 */

public class JWKBasedJwtToMapAdapter implements Decoder<Map<String, Object>, String> {
    // Issuer JWKS endpoint
    protected final String jwksUrl;
    // True if we dont want to actually verify the token
    protected final boolean skipVerification;
    // Logger
    Logger log = LoggerFactory.getLogger(this.getClass());
    /**
     * Constructor
     * @param jwksUrl  - the issuer endpoint where jwks is returned
     * @param skipVerification - you can choose to verify or not.
     */
    public JWKBasedJwtToMapAdapter(String jwksUrl, boolean skipVerification) {
        this.jwksUrl = jwksUrl;
        this.skipVerification = skipVerification;
    }

    /**
     * Decode and optionally verify the token
     * @param token  jwt token to verify
     * @return  the claims
     * @throws DecodeException if there is something wrong - invalid token etc
     */
    public Map<String, Object> decode(String token) throws DecodeException {
        var jwt = JWT.decode(token);
        JwtClaims verifiedClaims;
        var clientId = jwt.getClaim("aud").asString();

        try {
            // OK look up jwks and verify
            log.trace("Using key:" + jwt.getKeyId() + " to search: " + jwksUrl);
            var httpsJkws = new HttpsJwks(jwksUrl);
            var httpsJwksKeyResolver = new HttpsJwksVerificationKeyResolver(httpsJkws);
            var jwtConsumer = new JwtConsumerBuilder()
                    .setVerificationKeyResolver(httpsJwksKeyResolver)
                    .setExpectedAudience(clientId)
                    .build();
            if (skipVerification) { // Dont verify
                jwtConsumer.setSkipVerificationKeyResolutionOnNone(true);
            }
            verifiedClaims = jwtConsumer.processToClaims(token);
        } catch (InvalidJwtException e) {
            throw new DecodeException();
        }
        // OK were here ... so valid
        return verifiedClaims.getClaimsMap();
    }
}
