package com.atlassian.bitbucket.mesh.auth;

import com.atlassian.bitbucket.mesh.util.KeyUtils;

import javax.annotation.Nonnull;
import java.security.PublicKey;
import java.util.Collection;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

import static java.util.Objects.requireNonNull;

public class DefaultPublicKeyRegistry implements PublicKeyRegistry {

    private final ConcurrentMap<String, ConcurrentMap<String, PublicKey>> keysByNode;

    public DefaultPublicKeyRegistry() {
        keysByNode = new ConcurrentHashMap<>();
    }

    @Nonnull
    @Override
    public String add(@Nonnull String nodeId, @Nonnull PublicKey publicKey) {
        requireNonNull(nodeId, "nodeId");
        requireNonNull(publicKey, "publicKey");

        String fingerprint = KeyUtils.getFingerprint(publicKey);
        keysByNode.computeIfAbsent(nodeId, ignored -> new ConcurrentHashMap<>())
                .put(fingerprint, publicKey);

        return fingerprint;
    }

    @Override
    public PublicKey findByFingerprint(@Nonnull String nodeId, @Nonnull String fingerprint) {
        requireNonNull(nodeId, "nodeId");
        requireNonNull(fingerprint, "fingerprint");

        ConcurrentMap<String, PublicKey> keysByFingerprint = keysByNode.get(nodeId);
        if (keysByFingerprint == null) {
            return null;
        }
        return keysByFingerprint.get(fingerprint);
    }

    @Override
    public boolean isRegistered(@Nonnull String nodeId) {
        requireNonNull(nodeId, "nodeId");

        ConcurrentMap<String, PublicKey> keysByFingerprint = keysByNode.get(nodeId);
        return !(keysByFingerprint == null || keysByFingerprint.isEmpty());
    }

    @Override
    public boolean removeAll(@Nonnull String nodeId) {
        requireNonNull(nodeId, "nodeId");

        ConcurrentMap<String, PublicKey> keysByFingerprint = keysByNode.remove(nodeId);
        return !(keysByFingerprint == null || keysByFingerprint.isEmpty());
    }

    @Override
    public boolean remove(@Nonnull String nodeId, @Nonnull PublicKey publicKey) {
        requireNonNull(nodeId, "nodeId");
        requireNonNull(publicKey, "publicKey");

        //This doesn't use computeIfAbsent because there's no reason to add a fingerprint map if there isn't one
        //when we're trying to remove a key
        ConcurrentMap<String, PublicKey> keysByFingerprint = keysByNode.get(nodeId);
        if (keysByFingerprint == null) {
            return false;
        }

        return keysByFingerprint.remove(KeyUtils.getFingerprint(publicKey), publicKey);
    }

    @Override
    public void set(@Nonnull String nodeId, @Nonnull Collection<PublicKey> publicKeys) {
        requireNonNull(nodeId, "nodeId");

        if (requireNonNull(publicKeys, "publicKeys").isEmpty()) {
            removeAll(nodeId);
        } else {
            ConcurrentMap<String, PublicKey> newKeys = new ConcurrentHashMap<>();
            for (PublicKey publicKey : publicKeys) {
                newKeys.put(KeyUtils.getFingerprint(publicKey), publicKey);
            }

            keysByNode.put(nodeId, newKeys);
        }
    }
}
