public class GCPAuthManagementService{

    // Class to deserialise response from the Google token URL and get the
    // ID Token
    public class GoogleAuthResponse {
        public String id_token;
    }
    
    
    /**
     * Get GCP service account keys from metadata
     * @return the metadata record for the GCP service account
     */
    public GCP_Key__mdt getGCPServiceAccount(String functionName) {
        List<GCP_Key__mdt> gcp_sa_keys = new List<GCP_Key__mdt>();
        gcp_sa_keys = [SELECT
            client_email__c,
            audience__c,
            token_endpoint__c,
            query_fields__c,
            object__c
            FROM GCP_Key__mdt
            WHERE function_name__c =: functionName];
            
        System.debug('getGCPServiceAccount: ' + gcp_sa_keys);    

        if (gcp_sa_keys.size() > 0) {
            return gcp_sa_keys[0];
        } else {
            throw new GCPException('Cannot find GCP Service Account Keys');
        }
    }
    
    /**
     * Constructs the JWT and invokes the Google Auth endpoint
     * @param serviceAccount has the GCP keys obtained from the
     * custom metadata object
     * @return the auth response from GCP containing the id token
     */
    public GoogleAuthResponse getGCPAuthToken(GCP_Key__mdt serviceAccount) {

        GoogleAuthResponse result = new GoogleAuthResponse();

        Auth.JWT jwt = new Auth.JWT();
        jwt.setAud(serviceAccount.token_endpoint__c);
        jwt.setIss(serviceAccount.client_email__c);
        
        System.debug('getGCPAuthToken jwt: ' + jwt);

        // Additional claims to set scope
        Map<String, Object> claims = new Map<String, Object>();
        claims.put('target_audience', serviceAccount.audience__c);
        jwt.setAdditionalClaims(claims);

        // Create the object that signs the JWT bearer token
        Auth.JWS jws = new Auth.JWS(jwt, 'google_cloud');

        // Get the resulting JWS in case debugging is required
        String token = jws.getCompactSerialization();

        // Set the token endpoint that the JWT bearer token is posted to
        String tokenEndpoint = serviceAccount.token_endpoint__c;

        Auth.JWTBearerTokenExchange bearer =
                new Auth.JWTBearerTokenExchange(tokenEndpoint, jws);

        if (!Test.isRunningTest()) {
            System.HttpResponse response = bearer.getHTTPResponse();

            result =
                (GoogleAuthResponse) JSON.deserialize(
                    response.getBody(), GoogleAuthResponse.class);

        System.debug('GoogleAuthResponse result: ' + result);
                    
                    
        } else {
            result.id_token = 'IN TEST';
        }

        return result;
    }    
    
    
    public class GCPException extends Exception {}    

}