package com.atlassian.bitbucket.mesh.auth;

import com.atlassian.bitbucket.mesh.event.EventListener;
import com.atlassian.bitbucket.mesh.execution.ExecutionLocal;
import com.atlassian.bitbucket.mesh.execution.ExecutionManager;
import com.atlassian.bitbucket.mesh.git.exception.FailedPreconditionException;
import com.atlassian.bitbucket.mesh.network.MeshNodeRemovedEvent;
import com.atlassian.bitbucket.mesh.util.KeyUtils;
import com.google.common.collect.Collections2;
import io.jsonwebtoken.SignatureAlgorithm;
import io.jsonwebtoken.security.Keys;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.TaskScheduler;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.nio.file.Path;
import java.security.KeyPair;
import java.security.PublicKey;
import java.time.Duration;
import java.time.Instant;
import java.util.Map;
import java.util.Optional;
import java.util.Set;

import static com.atlassian.bitbucket.mesh.rpc.RpcConstants.Token.NID_CONTROL_PLANE;
import static com.atlassian.bitbucket.mesh.rpc.RpcConstants.Token.NID_SIDECAR;
import static java.util.Objects.requireNonNull;

public class DefaultAuthenticationManager implements AuthenticationManager {

    static final String PEM_CONTROL_PLANE = "control-plane.pem";
    static final String PEM_SIGNING_KEY = "signing-key.pem";

    private static final Logger log = LoggerFactory.getLogger(DefaultAuthenticationManager.class);

    private final Path controlPlanePem;
    private final ExecutionManager executionManager;
    private final PublicKeyRegistry keyRegistry;
    private final ExecutionLocal<Principal> principalLocal;
    private final Path signingKeyPem;
    private final TaskScheduler taskScheduler;

    private volatile SigningKey signingKey;

    public DefaultAuthenticationManager(Path configDir, ExecutionManager executionManager,
                                        PublicKeyRegistry keyRegistry, TaskScheduler taskScheduler) {
        this.executionManager = executionManager;
        this.keyRegistry = keyRegistry;
        this.taskScheduler = taskScheduler;

        controlPlanePem = configDir.resolve(PEM_CONTROL_PLANE);
        principalLocal = executionManager.newExecutionLocal();
        signingKeyPem = configDir.resolve(PEM_SIGNING_KEY);
    }

    @Nonnull
    @Override
    public Optional<Principal> getPrincipal() {
        return Optional.ofNullable(principalLocal.get());
    }

    @Nonnull
    @Override
    public Principal getPrincipalOrThrow() {
        return getPrincipal().orElseThrow(() -> new FailedPreconditionException("No Principal is authenticated"));
    }

    @Nonnull
    @Override
    public SigningKey getSigningKey() {
        return signingKey;
    }

    @Override
    public boolean isProvisioned() {
        return keyRegistry.isRegistered(NID_CONTROL_PLANE);
    }

    @EventListener
    public void onMeshNodeRemoved(MeshNodeRemovedEvent event) {
        String nodeId = event.getNode().getId();
        if (keyRegistry.removeAll(nodeId)) {
            log.debug("{}: Removed all registered public keys", nodeId);
        }
    }

    public void prepareKeys() throws IOException {
        PublicKey controlKey = KeyUtils.readPublicKey(controlPlanePem);
        if (controlKey == null) {
            log.info("No control plane key has been provisioned");
        } else {
            String fingerprint = keyRegistry.add(NID_CONTROL_PLANE, controlKey);
            log.debug("Registered provisioned control plane key (Fingerprint: {})", fingerprint);
        }

        KeyPair nodeKey = KeyUtils.readKeyPair(signingKeyPem);
        boolean newKey = nodeKey == null;
        if (newKey) {
            nodeKey = Keys.keyPairFor(SignatureAlgorithm.RS256);
            KeyUtils.writeKeyPair(signingKeyPem, nodeKey);
        }

        signingKey = new SimpleSigningKey(nodeKey);
        if (newKey) {
            log.info("Generated signing key (Fingerprint: {})", signingKey.getFingerprint());
        } else {
            log.debug("Loaded signing key (Fingerprint: {})", signingKey.getFingerprint());
        }
    }

    @Nonnull
    @Override
    public PublicKey provision(@Nonnull PublicKey publicKey) {
        requireNonNull(publicKey, "publicKey");

        Principal principal = principalLocal.get();
        if (principal != null && principal.getRole() != ApplicationRole.CONTROL_PLANE) {
            throw new AuthorisationException("Only the control plane can provision this node");
        }

        //Write the control plane's public key to disk first and, if that succeeds, add it to the registry
        //so it can be resolved to verify tokens for subsequent requests
        writeControlPlaneKey(publicKey);
        keyRegistry.add(NID_CONTROL_PLANE, publicKey);

        return signingKey.getPublic();
    }

    @Override
    public void updatePublicKeys(@Nonnull UpdatePublicKeysParameters parameters) {
        requireNonNull(parameters, "parameters");

        Set<ExpirablePublicKey> controlPlaneKeys = parameters.getControlPlaneKeys();
        if (controlPlaneKeys.isEmpty()) {
            log.debug("No updated control plane keys were provided");
        } else {
            //If any control plane keys were provided, look for one that doesn't expire. If one is found,
            //write that key to disk so it's available after the next restart
            controlPlaneKeys.stream()
                    .filter(key -> !key.isExpirable())
                    .findFirst()
                    .ifPresent(key -> writeControlPlaneKey(key.getPublicKey()));

            //After the control plane's key is successfully persisted to disk, then update the registry
            registerKeys(NID_CONTROL_PLANE, controlPlaneKeys, parameters.isComplete());
        }

        Set<ExpirablePublicKey> sidecarKeys = parameters.getSidecarKeys();
        if (sidecarKeys.isEmpty()) {
            log.debug("No updated sidecar keys were provided");
        } else {
            registerKeys(NID_SIDECAR, sidecarKeys, parameters.isComplete());
        }

        Map<String, Set<ExpirablePublicKey>> keysByNode = parameters.getKeysByNode();
        if (keysByNode.isEmpty()) {
            log.debug("No updated node keys were provided");
        } else {
            keysByNode.forEach((nodeId, nodeKeys) -> registerKeys(nodeId, nodeKeys, parameters.isComplete()));
        }
    }

    @Override
    public <T> T withPrincipal(@Nonnull Principal principal, @Nonnull PrincipalCallback<T> callback) {
        requireNonNull(principal, "principal");
        requireNonNull(callback, "callback");

        if (executionManager.getContext().isPresent()) {
            //Attach the principal to the current ExecutionContext. Note that, for this path, the
            //principal will _intentionally_ still be active after this method returns
            return callWithPrincipal(principal, callback);
        }

        //Otherwise, create an ExecutionContext and associate the principal with that. The context
        //will automatically be torn down after the PrincipalCallback is finished
        return executionManager.withContext(ignored -> callWithPrincipal(principal, callback));
    }

    private <T> T callWithPrincipal(Principal principal, PrincipalCallback<T> callback) {
        principalLocal.set(principal);

        return callback.callWithPrincipal(principal);
    }

    private void registerKeys(String nodeId, Set<ExpirablePublicKey> keys, boolean replace) {
        if (replace) {
            keyRegistry.set(nodeId, Collections2.transform(keys, ExpirablePublicKey::getPublicKey));
        } else {
            keys.stream()
                    .map(ExpirablePublicKey::getPublicKey)
                    .forEach(publicKey -> keyRegistry.add(nodeId, publicKey));
        }

        scheduleExpiryTasks(nodeId, keys);
    }

    private void scheduleExpiryTasks(String nodeId, Set<ExpirablePublicKey> keys) {
        for (ExpirablePublicKey key : keys) {
            Duration interval = key.getExpiryInterval();
            if (interval == null) {
                continue;
            }

            taskScheduler.schedule(() -> keyRegistry.remove(nodeId, key.getPublicKey()), Instant.now().plus(interval));
        }
    }

    private void writeControlPlaneKey(PublicKey publicKey) {
        try {
            KeyUtils.writePublicKey(controlPlanePem, publicKey);
            log.info("Registered new control plane key (Fingerprint: {})", KeyUtils.getFingerprint(publicKey));
        } catch (IOException e) {
            throw new FailedPreconditionException("Failed to write the control plane's public key", e);
        }
    }
}
