/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.kafka.common.security.oauthbearer.internals.secured;

import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.common.utils.Time;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;

import org.jose4j.jwk.PublicJsonWebKey;
import org.jose4j.jws.JsonWebSignature;
import org.jose4j.jwt.ReservedClaimNames;
import org.jose4j.lang.JoseException;

import java.io.IOException;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

public class AccessTokenBuilder {

    private final String scopeClaimName = "scope";

    private final Long issuedAtSeconds;

    private final ObjectMapper objectMapper = new ObjectMapper();

    private String alg;

    private String audience;

    private String subject = "jdoe";

    private String subjectClaimName = ReservedClaimNames.SUBJECT;

    private Object scope = "engineering";

    private Long expirationSeconds;

    private PublicJsonWebKey jwk;

    private final Map<String, String> customClaims = new HashMap<>();

    public AccessTokenBuilder() {
        this(new MockTime());
    }

    public AccessTokenBuilder(Time time) {
        this.issuedAtSeconds = time.milliseconds() / 1000;
        this.expirationSeconds = this.issuedAtSeconds + 60;
    }

    public String alg() {
        return alg;
    }

    public AccessTokenBuilder alg(String alg) {
        this.alg = alg;
        return this;
    }

    public AccessTokenBuilder audience(String audience) {
        this.audience = audience;
        return this;
    }

    public String subject() {
        return subject;
    }

    public AccessTokenBuilder subject(String subject) {
        this.subject = subject;
        return this;
    }

    public String subjectClaimName() {
        return subjectClaimName;
    }

    public AccessTokenBuilder subjectClaimName(String subjectClaimName) {
        this.subjectClaimName = subjectClaimName;
        return this;
    }

    public Object scope() {
        return scope;
    }

    public AccessTokenBuilder scope(Object scope) {
        this.scope = scope;

        if (scope instanceof String) {
            return this;
        } else if (scope instanceof Collection) {
            return this;
        } else {
            throw new IllegalArgumentException(String.format("%s parameter must be a %s or a %s containing %s",
                scopeClaimName,
                String.class.getName(),
                Collection.class.getName(),
                String.class.getName()));
        }
    }

    public String scopeClaimName() {
        return scopeClaimName;
    }

    public Long issuedAtSeconds() {
        return issuedAtSeconds;
    }

    public Long expirationSeconds() {
        return expirationSeconds;
    }

    public AccessTokenBuilder expirationSeconds(Long expirationSeconds) {
        this.expirationSeconds = expirationSeconds;
        return this;
    }

    public PublicJsonWebKey jwk() {
        return jwk;
    }

    public AccessTokenBuilder jwk(PublicJsonWebKey jwk) {
        this.jwk = jwk;
        return this;
    }

    public AccessTokenBuilder addCustomClaim(String name, String value) {
        String validatedName = ClaimValidationUtils.validateClaimNameOverride("claim name", name);
        String validatedValue = ClaimValidationUtils.validateClaimNameOverride(validatedName, value);

        customClaims.put(validatedName, validatedValue);
        return this;
    }

    @SuppressWarnings("unchecked")
    public String build() throws JoseException, IOException {
        ObjectNode node = objectMapper.createObjectNode();

        if (audience != null)
            node.put(ReservedClaimNames.AUDIENCE, audience);

        if (subject != null)
            node.put(subjectClaimName, subject);

        if (scope instanceof String) {
            node.put(scopeClaimName, (String) scope);
        } else if (scope instanceof Collection) {
            ArrayNode child = node.putArray(scopeClaimName);
            ((Collection<String>) scope).forEach(child::add);
        } else {
            throw new IllegalArgumentException(String.format("%s claim must be a %s or a %s containing %s",
                scopeClaimName,
                String.class.getName(),
                Collection.class.getName(),
                String.class.getName()));
        }

        if (issuedAtSeconds != null)
            node.put(ReservedClaimNames.ISSUED_AT, issuedAtSeconds);

        if (expirationSeconds != null)
            node.put(ReservedClaimNames.EXPIRATION_TIME, expirationSeconds);

        for (Map.Entry<String, String> claim : customClaims.entrySet()) {
            node.put(claim.getKey(), claim.getValue());
        }

        String json = objectMapper.writeValueAsString(node);

        JsonWebSignature jws = new JsonWebSignature();
        jws.setPayload(json);

        if (jwk != null) {
            jws.setKey(jwk.getPrivateKey());
            jws.setKeyIdHeaderValue(jwk.getKeyId());
        }

        if (alg != null)
            jws.setAlgorithmHeaderValue(alg);

        return jws.getCompactSerialization();
    }

}
