package com.github.scribejava.apis;

import java.net.MalformedURLException;
import java.net.URL;
import java.text.ParseException;
import java.util.Map;
import java.util.Optional;
import java.util.Set;

import com.nimbusds.jose.JOSEException;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.jwk.source.JWKSource;
import com.nimbusds.jose.jwk.source.RemoteJWKSet;
import com.nimbusds.jose.proc.BadJOSEException;
import com.nimbusds.jose.proc.JWSKeySelector;
import com.nimbusds.jose.proc.JWSVerificationKeySelector;
import com.nimbusds.jose.proc.SecurityContext;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.proc.ConfigurableJWTProcessor;
import com.nimbusds.jwt.proc.DefaultJWTClaimsVerifier;
import com.nimbusds.jwt.proc.DefaultJWTProcessor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class GoogleTokenVerifier {
    private static final Logger logger = LoggerFactory.getLogger("JWT");

    public static Optional<String> validateToken(String clientId, String idToken) {

        // Create a JWT processor for the access tokens
        ConfigurableJWTProcessor<SecurityContext> jwtProcessor = new DefaultJWTProcessor<>();

        // The public RSA keys to validate the signatures will be sourced from the
        // OAuth 2.0 server's JWK set, published at a well-known URL. The RemoteJWKSet
        // object caches the retrieved keys to speed up subsequent look-ups and can
        // also handle key-rollover
        JWKSource<SecurityContext> keySource =
                null;
        try {
            keySource = new RemoteJWKSet<>(new URL("https://www.googleapis.com/oauth2/v3/certs"));
        } catch (MalformedURLException e) {
            return Optional.empty();
        }

        // The expected JWS algorithm of the access tokens (agreed out-of-band)
        JWSAlgorithm expectedJWSAlg = JWSAlgorithm.RS256;

        // Configure the JWT processor with a key selector to feed matching public
        // RSA keys sourced from the JWK set URL
        JWSKeySelector<SecurityContext> keySelector =
                new JWSVerificationKeySelector<>(expectedJWSAlg, keySource);

        jwtProcessor.setJWSKeySelector(keySelector);

        // Set the required JWT claims for access tokens issued by the server
        jwtProcessor.setJWTClaimsSetVerifier(new DefaultJWTClaimsVerifier<>(
                new JWTClaimsSet.Builder()
                        .issuer("https://accounts.google.com")
                        .audience(clientId)
                        .build(),
                Set.of("exp", "iat", "aud", "email")));

        // Process the token
        Map<String, Object> claimsSet;
        try {
            claimsSet = jwtProcessor.process(idToken, null).toJSONObject();
        } catch (BadJOSEException | JOSEException | ParseException e) {
            logger.error(e.getMessage(), e);
            return Optional.empty();
        }

        String email = (String) claimsSet.get("email");
        boolean verified = claimsSet.get("email_verified").equals(true);
        return verified ? Optional.of(email) : Optional.empty();
    }
}