package com.atlassian.bitbucket.internal.mirroring.mirror.auth;

import com.atlassian.bitbucket.i18n.I18nService;
import com.atlassian.bitbucket.internal.mirroring.mirror.MirrorKeyUploadFailedException;
import com.atlassian.bitbucket.internal.mirroring.mirror.client.InternalUpstreamClientFactory;
import com.atlassian.bitbucket.internal.mirroring.mirror.dao.AoUpstreamServer;
import com.atlassian.bitbucket.internal.mirroring.mirror.dao.UpstreamServerDao;
import com.atlassian.bitbucket.internal.mirroring.mirror.ssh.MirrorKeyStorageException;
import com.atlassian.bitbucket.internal.mirroring.mirror.ssh.SshKeyStore;
import com.atlassian.bitbucket.mirroring.mirror.NoSuchUpstreamException;
import com.atlassian.bitbucket.mirroring.mirror.UpstreamServer;
import com.atlassian.bitbucket.mirroring.mirror.UpstreamServerType;
import com.atlassian.bitbucket.scm.mirror.MirrorSyncCommandParameters;
import com.atlassian.bitbucket.server.ApplicationPropertiesService;
import com.atlassian.fugue.Pair;
import com.atlassian.sal.api.transaction.TransactionTemplate;
import com.atlassian.util.concurrent.Promise;
import com.google.common.base.MoreObjects;
import com.google.common.collect.Maps;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.annotation.Nonnull;
import java.io.File;
import java.net.URI;
import java.util.Optional;
import java.util.concurrent.ConcurrentMap;

import static com.atlassian.bitbucket.internal.mirroring.mirror.MirrorDescriptionUtils.describe;
import static com.atlassian.util.concurrent.Promises.promise;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;

/**
 * Manages credentials used by mirror servers for fetching from the upstream server. If the upstream is a Bitbucket
 * Server instance, this will be an SSH key. If the upstream server is a Bitbucket Cloud instance, the credential will
 * be an OAUTH token.
 * <p>
 * The manager is responsible for registering new credentials (generating and uploading SSH key for Server, or
 * obtaining an OAUTH access token for Cloud), renewing credentials when they expire and cleaning up credentials when
 * the mirror is uninstalled from the upstream.
 * <p>
 * This class ensures that for any upstream, there is only ever one ongoing operation that mutates the credentials at
 * any one time. Any client who, by requesting or refreshing credentials, attempts to start a new update while one is
 * already ongoing, will receive the promise for the already-existing operation. This ensures two (or more)
 * simultaneously running fetch tasks do not overwrite keys locally or on the upstream but instead each waits on the
 * result of the update initiated by the first.
 */
@Component
public class DefaultSyncCredentialsManager implements SyncCredentialsManager {

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

    private static final int NO_VERSION = -1;

    private final ConcurrentMap<String, Pair<Integer, Promise<CloudSyncCredentials>>> cloudCredentials;
    private final ConcurrentMap<String, Pair<Integer, Promise<ServerSyncCredentials>>> serverCredentials;
    private final I18nService i18nService;
    private final ApplicationPropertiesService propertiesService;
    private final SshKeyStore sshKeyStore;
    private final TransactionTemplate transactionTemplate;
    private final InternalUpstreamClientFactory upstreamClientFactory;
    private final UpstreamServerDao upstreamServerDao;

    @Autowired
    public DefaultSyncCredentialsManager(I18nService i18nService,
                                         ApplicationPropertiesService propertiesService,
                                         SshKeyStore sshKeyStore,
                                         TransactionTemplate transactionTemplate,
                                         InternalUpstreamClientFactory upstreamClientFactory,
                                         UpstreamServerDao upstreamServerDao) {

        this.i18nService = i18nService;
        this.propertiesService = propertiesService;
        this.sshKeyStore = sshKeyStore;
        this.transactionTemplate = transactionTemplate;
        this.upstreamClientFactory = upstreamClientFactory;
        this.upstreamServerDao = upstreamServerDao;

        cloudCredentials = Maps.newConcurrentMap();
        serverCredentials = Maps.newConcurrentMap();
    }

    @Nonnull
    @Override
    public Promise<SyncCredentials> getCredentials(@Nonnull String upstreamId) {
        return computeCredentials(NO_VERSION, getUpstreamOrFail(upstreamId));
    }

    @SuppressWarnings("unchecked")
    @Nonnull
    @Override
    public Promise<SyncCredentials> refresh(@Nonnull SyncCredentials credentials) {
        return (Promise) computeCredentials(
                ((BaseSyncCredentials) credentials).getVersion(),
                ((BaseSyncCredentials) credentials).getUpstream());
    }

    @SuppressWarnings("unchecked")
    @Nonnull
    private Promise<SyncCredentials> computeCredentials(int callerCredentialsVersion, @Nonnull UpstreamServer upstream) {
        switch (upstream.getType()) {
            case BITBUCKET_CLOUD:
                return (Promise) computeCredentialsFor(
                        cloudCredentials, callerCredentialsVersion, this::getCloudCredentials, upstream);
            case BITBUCKET_SERVER:
                return (Promise) computeCredentialsFor(
                        serverCredentials, callerCredentialsVersion, this::getServerCredentials, upstream);
            default:
                throw new IllegalArgumentException("Upstream " + upstream.getId() + " of type " + upstream.getType() +
                        " is not supported");
        }
    }

    /**
     * Returns a {@link Promise delayed computation} of credentials for the {@link UpstreamServer upstream}.
     * <p>
     * If credentials have not yet been retrieved a new promise of the credentials is returned that will retrieve/store
     * the credentials upstream as implemented by {@code compute} (note that his refreshed computation may fail).
     * <p>
     * If credentials have already been retrieved/stored successfully and {@code refresh} is false, the completed
     * promise of the credentials is returned.
     * <p>
     * If credentials have already been retrieved/stored successfully and {@code refresh} is true, a new promise of the
     * credentials is returned that will retrieve/store the credentials upstream as defined by {@code compute}
     * (note that his refreshed computation may fail).
     * <p>
     * If an attempt to retrieve/store credentials have already been made but failed with error, a new promise of the
     * credentials is returned that will return the result of retrieving the credentials upstream as defined by
     * {@code compute} (note that this computation may fail, too).
     * <p>
     * Note: This method ensures that for the given {@link ConcurrentMap concurrent map} of credential promises there is
     * never more than one computation ongoing per upstream.
     *
     * @param <C> the type of credentials for the upstream server
     * @param allCredentials the {@link ConcurrentMap concurrent map} of all computed credentials for all upstreams
     * @param callerCredentialsVersion the {@link BaseSyncCredentials#getVersion() version} of the credentials the caller has
     * @param credentialsProvider the computation used to retrieve the credentials for the supplied upstream
     * @param upstream the {@link UpstreamServer upstream server}
     * @return a possibly already completed promise of the computed credentials
     */
    @SuppressWarnings("unchecked")
    private <C extends SyncCredentials> Promise<C> computeCredentialsFor(
            @Nonnull ConcurrentMap<String, Pair<Integer, Promise<C>>> allCredentials, int callerCredentialsVersion,
            @Nonnull CredentialsProvider<C> credentialsProvider, @Nonnull UpstreamServer upstream) {

        return (Promise) allCredentials.compute(upstream.getId(), (id, existing) -> {
            boolean firstComputationRequest = existing == null;
            boolean refreshRequested = callerCredentialsVersion != NO_VERSION;
            boolean inProgress = !firstComputationRequest && !existing.right().isDone();
            boolean inError = !firstComputationRequest && !inProgress && isError(existing.right());
            boolean callerCredentialsStale = !firstComputationRequest && existing.left() > callerCredentialsVersion;

            checkState(!firstComputationRequest || callerCredentialsVersion == NO_VERSION,
                    "Unexpected state: computation does not exist in the map but a refresh has been requested");

            if (inProgress || (callerCredentialsStale && !inError)) {
                // return the existing credentials computation if:
                // 1. the credential computation is ongoing (only one ongoing computation allowed at any one time per upstream)
                // 2. the caller has held on to stale credentials and the last computation was in error (upload or key
                //    storage failed)
                if (log.isTraceEnabled()) {
                    String reason = inProgress
                            ? "it is still ongoing"
                            : "the caller has stale credentials";
                    log.trace("Using existing credentials computation for {} because: {}", describe(upstream), reason);
                }

                return existing;
            }

            // (re-)compute if:
            // 1. this is the first computation for this server's lifetime
            // 2. the last computation was in error (upload or key storage failed)
            // 3. the last computation succeeded but the caller wants it refreshed

            // only mandate new credentials if refreshing or the previous computation failed. first requests
            boolean requiresNewCredentials = refreshRequested || inError;

            if (log.isTraceEnabled()) {
                String reason = firstComputationRequest
                        ? "this is the first computation for that upstream since the application started"
                        : refreshRequested
                            ? "the caller has asked for a refresh of the credentials"
                            : "the last computation resulted in an error";
                log.trace("Computing credentials for {} because: {}. {}", describe(upstream), reason);
            }

            return Pair.pair(callerCredentialsVersion + 1,
                    credentialsProvider.get(callerCredentialsVersion + 1, requiresNewCredentials, upstream));
        }).right();
    }

    private static Throwable detailsIfDebug(Throwable t) {
        return log.isDebugEnabled() ? t : null;
    }

    private <C> boolean isError(Promise<C> existing) {
        if (existing.isDone()) {
            try {
                existing.claim();
            } catch (Throwable e) {
                // it's unlikely an Error here would be because of the claim() call itself but rather from the
                // computation returned for the claim and the caller will get a repeat performance once they claim it so
                // it's generally OK to catch Throwable here
                return true;
            }
        }

        return false;
    }

    private Promise<File> createAndRegisterSshKeyFor(UpstreamServer upstream) {
        Pair<File, String> keyPair = sshKeyStore.generateKeyPair(upstream.getId(), getLabel());

        return upstreamClientFactory.createForServer(upstream).uploadSshKey(keyPair.right())
                .map(input -> keyPair.left())
                .done(files -> log.info("The SSH key for this mirror was uploaded to {}", describe(upstream)))
                .fail(t -> {
                    if (t instanceof MirrorKeyUploadFailedException) {
                        log.error("Uploading the SSH key for {} failed. Upstream fetches will not proceed until this is corrected.",
                                describe(upstream), detailsIfDebug(t));
                    } else if (t instanceof MirrorKeyStorageException) {
                        log.error("Storing the SSH key for {} failed. Upstream fetches will not proceed.",
                                describe(upstream), detailsIfDebug(t));
                    } else {
                        log.error("An unexpected error was encountered while uploading and storing the SSH key for {}. " +
                                "Synchronization will not proceed.", describe(upstream), detailsIfDebug(t));
                    }
                });
    }

    private boolean deleteSshKeyFor(@Nonnull UpstreamServer upstream) {
        deleteUpstreamSshKeys(upstream);
        return sshKeyStore.delete(upstream.getId());
    }

    private void deleteUpstreamSshKeys(UpstreamServer upstream) {
        try {
            upstreamClientFactory.createForServer(upstream).deleteSshKeys();
        } catch (RuntimeException e) {
            log.info("Failed to delete SSH key with upstream {}: '{}'",
                    upstream.getBaseUrl(), e.getMessage(), detailsIfDebug(e));
        }
    }

    private Promise<CloudSyncCredentials> getCloudCredentials(int nextVersion, boolean requiresNew, UpstreamServer upstream) {
        return upstreamClientFactory.createForCloud(upstream)
                .getAccessToken()
                .map(token -> new CloudSyncCredentials(token, upstream, nextVersion));
    }

    private String getLabel() {
        URI baseUrl = propertiesService.getBaseUrl();
        if (baseUrl != null) {
            return "mirror " + baseUrl.toASCIIString();
        }
        return "mirror " + propertiesService.getDisplayName() + "(" + propertiesService.getServerId() + ")";
    }

    private Promise<ServerSyncCredentials> getServerCredentials(int nextVersion, boolean requiresNew, UpstreamServer upstream) {
        if (!requiresNew) {
            Optional<File> keyFile = sshKeyStore.getPrivateKey(upstream.getId());
            if (keyFile.isPresent()) {
                log.trace("{}: SSH key found at {}", describe(upstream), keyFile.get().getAbsolutePath());
                return promise(new ServerSyncCredentials(keyFile.get(), upstream, nextVersion));
            }
        }

        if (requiresNew) {
            deleteSshKeyFor(upstream);
        }

        log.info("Generating and configuring an SSH key for {}", describe(upstream));
        return createAndRegisterSshKeyFor(upstream)
                .map(newFile -> new ServerSyncCredentials(newFile, upstream, nextVersion));
    }

    private AoUpstreamServer getUpstreamOrFail(@Nonnull String upstreamId) {
        return transactionTemplate.execute(() -> {
            AoUpstreamServer upstream = upstreamServerDao.getById(upstreamId);
            if (upstream == null) {
                throw new NoSuchUpstreamException(i18nService.createKeyedMessage("bitbucket.mirroring.no.such.upstream.server", upstreamId));
            }
            return upstream;
        });
    }

    private class CloudSyncCredentials extends BaseSyncCredentials {

        private final OAuthAccessToken accessToken;

        private CloudSyncCredentials(@Nonnull OAuthAccessToken accessToken, @Nonnull UpstreamServer upstream, int version) {
            super(upstream, version);

            this.accessToken = accessToken;
            checkArgument(upstream.getType() == UpstreamServerType.BITBUCKET_CLOUD, "Must be a cloud upstream");
        }

        @Nonnull
        @Override
        public MirrorSyncCommandParameters.Builder configure(@Nonnull MirrorSyncCommandParameters.Builder builder) {
            return builder
                    .username("x-token-auth")
                    .password(accessToken.getValue());
        }

        @Override
        public String toString() {
            return MoreObjects.toStringHelper(this)
                    .add("accessToken", accessToken)
                    .add("upstream", upstream)
                    .add("version", version)
                    .toString();
        }
    }

    private class ServerSyncCredentials extends BaseSyncCredentials {
        private final File keyFile;
        private final UpstreamServer upstream;

        private ServerSyncCredentials(@Nonnull File keyFile, @Nonnull UpstreamServer upstream, int version) {
            super(upstream, version);

            this.keyFile = keyFile;
            this.upstream = upstream;
        }

        @Nonnull
        @Override
        public MirrorSyncCommandParameters.Builder configure(@Nonnull MirrorSyncCommandParameters.Builder builder) {
            return builder.privateKey(keyFile);
        }

        @Override
        public String toString() {
            return MoreObjects.toStringHelper(this)
                    .add("keyFile", keyFile)
                    .add("upstream", upstream)
                    .add("version", version)
                    .toString();
        }
    }

    private abstract class BaseSyncCredentials implements SyncCredentials {
        protected final int version;
        protected final UpstreamServer upstream;

        private BaseSyncCredentials(@Nonnull UpstreamServer upstream, int version) {
            this.version = version;
            this.upstream = upstream;
        }


        private UpstreamServer getUpstream() {
            return upstream;
        }

        private int getVersion() {
            return version;
        }
    }

    @FunctionalInterface
    interface CredentialsProvider<C extends SyncCredentials> {
        /**
         * @return a {@link Promise} of the credentials for the supplied upstream
         */
        @Nonnull
        Promise<C> get(int nextVersion, boolean requiresNew, @Nonnull UpstreamServer upstream);
    }
}