/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License;
 * you may not use this file except in compliance with the Elastic License.
 */

package org.elasticsearch.xpack.ccr;

import org.elasticsearch.ElasticsearchStatusException;
import org.elasticsearch.action.Action;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.ActionRequest;
import org.elasticsearch.action.ActionResponse;
import org.elasticsearch.action.admin.cluster.state.ClusterStateRequest;
import org.elasticsearch.action.admin.cluster.state.ClusterStateResponse;
import org.elasticsearch.action.support.ContextPreservingActionListener;
import org.elasticsearch.action.admin.indices.stats.IndexShardStats;
import org.elasticsearch.action.admin.indices.stats.IndexStats;
import org.elasticsearch.action.admin.indices.stats.IndicesStatsRequest;
import org.elasticsearch.action.admin.indices.stats.IndicesStatsResponse;
import org.elasticsearch.action.admin.indices.stats.ShardStats;
import org.elasticsearch.client.Client;
import org.elasticsearch.client.FilterClient;
import org.elasticsearch.cluster.ClusterState;
import org.elasticsearch.cluster.metadata.IndexMetaData;
import org.elasticsearch.common.util.concurrent.ThreadContext;
import org.elasticsearch.common.CheckedConsumer;
import org.elasticsearch.index.IndexNotFoundException;
import org.elasticsearch.index.engine.CommitStats;
import org.elasticsearch.index.engine.Engine;
import org.elasticsearch.index.shard.ShardId;
import org.elasticsearch.license.RemoteClusterLicenseChecker;
import org.elasticsearch.license.XPackLicenseState;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.xpack.ccr.action.ShardFollowTask;
import org.elasticsearch.xpack.core.XPackPlugin;

import java.util.Collections;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;
import java.util.function.BiConsumer;
import java.util.function.BooleanSupplier;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;

/**
 * Encapsulates licensing checking for CCR.
 */
public final class CcrLicenseChecker {

    private final BooleanSupplier isCcrAllowed;

    /**
     * Constructs a CCR license checker with the default rule based on the license state for checking if CCR is allowed.
     */
    CcrLicenseChecker() {
        this(XPackPlugin.getSharedLicenseState()::isCcrAllowed);
    }

    /**
     * Constructs a CCR license checker with the specified boolean supplier.
     *
     * @param isCcrAllowed a boolean supplier that should return true if CCR is allowed and false otherwise
     */
    public CcrLicenseChecker(final BooleanSupplier isCcrAllowed) {
        this.isCcrAllowed = Objects.requireNonNull(isCcrAllowed);
    }

    /**
     * Returns whether or not CCR is allowed.
     *
     * @return true if CCR is allowed, otherwise false
     */
    public boolean isCcrAllowed() {
        return isCcrAllowed.getAsBoolean();
    }

    /**
     * Fetches the leader index metadata and history UUIDs for leader index shards from the remote cluster.
     * Before fetching the index metadata, the remote cluster is checked for license compatibility with CCR.
     * If the remote cluster is not licensed for CCR, the {@code onFailure} consumer is is invoked. Otherwise,
     * the specified consumer is invoked with the leader index metadata fetched from the remote cluster.
     *
     * @param client        the client
     * @param clusterAlias  the remote cluster alias
     * @param leaderIndex   the name of the leader index
     * @param onFailure     the failure consumer
     * @param consumer      the consumer for supplying the leader index metadata and historyUUIDs of all leader shards
     * @param <T>           the type of response the listener is waiting for
     */
    public <T> void checkRemoteClusterLicenseAndFetchLeaderIndexMetadataAndHistoryUUIDs(
            final Client client,
            final String clusterAlias,
            final String leaderIndex,
            final Consumer<Exception> onFailure,
            final BiConsumer<String[], IndexMetaData> consumer) {

        final ClusterStateRequest request = new ClusterStateRequest();
        request.clear();
        request.metaData(true);
        request.indices(leaderIndex);
        checkRemoteClusterLicenseAndFetchClusterState(
                client,
                Collections.emptyMap(),
                clusterAlias,
                request,
                onFailure,
                leaderClusterState -> {
                    IndexMetaData leaderIndexMetaData = leaderClusterState.getMetaData().index(leaderIndex);
                    if (leaderIndexMetaData == null) {
                        onFailure.accept(new IndexNotFoundException(leaderIndex));
                        return;
                    }

                    final Client leaderClient = client.getRemoteClusterClient(clusterAlias);
                    fetchLeaderHistoryUUIDs(leaderClient, leaderIndexMetaData, onFailure, historyUUIDs -> {
                        consumer.accept(historyUUIDs, leaderIndexMetaData);
                    });
                },
                licenseCheck -> indexMetadataNonCompliantRemoteLicense(leaderIndex, licenseCheck),
                e -> indexMetadataUnknownRemoteLicense(leaderIndex, clusterAlias, e));
    }

    /**
     * Fetches the leader cluster state from the remote cluster by the specified cluster state request. Before fetching the cluster state,
     * the remote cluster is checked for license compliance with CCR. If the remote cluster is not licensed for CCR,
     * the {@code onFailure} consumer is invoked. Otherwise, the specified consumer is invoked with the leader cluster state fetched from
     * the remote cluster.
     *
     * @param client                     the client
     * @param clusterAlias               the remote cluster alias
     * @param headers                    the headers to use for leader client
     * @param request                    the cluster state request
     * @param onFailure                  the failure consumer
     * @param leaderClusterStateConsumer the leader cluster state consumer
     * @param <T>                        the type of response the listener is waiting for
     */
    public <T> void checkRemoteClusterLicenseAndFetchClusterState(
            final Client client,
            final Map<String, String> headers,
            final String clusterAlias,
            final ClusterStateRequest request,
            final Consumer<Exception> onFailure,
            final Consumer<ClusterState> leaderClusterStateConsumer) {
        checkRemoteClusterLicenseAndFetchClusterState(
                client,
                headers,
                clusterAlias,
                request,
                onFailure,
                leaderClusterStateConsumer,
                CcrLicenseChecker::clusterStateNonCompliantRemoteLicense,
                e -> clusterStateUnknownRemoteLicense(clusterAlias, e));
    }

    /**
     * Fetches the leader cluster state from the remote cluster by the specified cluster state request. Before fetching the cluster state,
     * the remote cluster is checked for license compliance with CCR. If the remote cluster is not licensed for CCR,
     * the {@code onFailure} consumer is invoked. Otherwise, the specified consumer is invoked with the leader cluster state fetched from
     * the remote cluster.
     *
     * @param client                     the client
     * @param clusterAlias               the remote cluster alias
     * @param headers                    the headers to use for leader client
     * @param request                    the cluster state request
     * @param onFailure                  the failure consumer
     * @param leaderClusterStateConsumer the leader cluster state consumer
     * @param nonCompliantLicense        the supplier for when the license state of the remote cluster is non-compliant
     * @param unknownLicense             the supplier for when the license state of the remote cluster is unknown due to failure
     * @param <T>                        the type of response the listener is waiting for
     */
    private <T> void checkRemoteClusterLicenseAndFetchClusterState(
            final Client client,
            final Map<String, String> headers,
            final String clusterAlias,
            final ClusterStateRequest request,
            final Consumer<Exception> onFailure,
            final Consumer<ClusterState> leaderClusterStateConsumer,
            final Function<RemoteClusterLicenseChecker.LicenseCheck, ElasticsearchStatusException> nonCompliantLicense,
            final Function<Exception, ElasticsearchStatusException> unknownLicense) {
        // we have to check the license on the remote cluster
        new RemoteClusterLicenseChecker(client, XPackLicenseState::isCcrAllowedForOperationMode).checkRemoteClusterLicenses(
                Collections.singletonList(clusterAlias),
                new ActionListener<RemoteClusterLicenseChecker.LicenseCheck>() {

                    @Override
                    public void onResponse(final RemoteClusterLicenseChecker.LicenseCheck licenseCheck) {
                        if (licenseCheck.isSuccess()) {
                            final Client leaderClient = wrapClient(client.getRemoteClusterClient(clusterAlias), headers);
                            final ActionListener<ClusterStateResponse> clusterStateListener =
                                    ActionListener.wrap(s -> leaderClusterStateConsumer.accept(s.getState()), onFailure);
                            // following an index in remote cluster, so use remote client to fetch leader index metadata
                            leaderClient.admin().cluster().state(request, clusterStateListener);
                        } else {
                            onFailure.accept(nonCompliantLicense.apply(licenseCheck));
                        }
                    }

                    @Override
                    public void onFailure(final Exception e) {
                        onFailure.accept(unknownLicense.apply(e));
                    }

                });
    }

    /**
     * Fetches the history UUIDs for leader index on per shard basis using the specified leaderClient.
     *
     * @param leaderClient                              the leader client
     * @param leaderIndexMetaData                       the leader index metadata
     * @param onFailure                                 the failure consumer
     * @param historyUUIDConsumer                       the leader index history uuid and consumer
     */
    // NOTE: Placed this method here; in order to avoid duplication of logic for fetching history UUIDs
    // in case of following a local or a remote cluster.
    public void fetchLeaderHistoryUUIDs(
        final Client leaderClient,
        final IndexMetaData leaderIndexMetaData,
        final Consumer<Exception> onFailure,
        final Consumer<String[]> historyUUIDConsumer) {

        String leaderIndex = leaderIndexMetaData.getIndex().getName();
        CheckedConsumer<IndicesStatsResponse, Exception> indicesStatsHandler = indicesStatsResponse -> {
            IndexStats indexStats = indicesStatsResponse.getIndices().get(leaderIndex);
            String[] historyUUIDs = new String[leaderIndexMetaData.getNumberOfShards()];
            for (IndexShardStats indexShardStats : indexStats) {
                for (ShardStats shardStats : indexShardStats) {
                    // Ignore replica shards as they may not have yet started and
                    // we just end up overwriting slots in historyUUIDs
                    if (shardStats.getShardRouting().primary() == false) {
                        continue;
                    }

                    CommitStats commitStats = shardStats.getCommitStats();
                    if (commitStats == null) {
                        onFailure.accept(new IllegalArgumentException("leader index's commit stats are missing"));
                        return;
                    }
                    String historyUUID = commitStats.getUserData().get(Engine.HISTORY_UUID_KEY);
                    ShardId shardId = shardStats.getShardRouting().shardId();
                    historyUUIDs[shardId.id()] = historyUUID;
                }
            }
            for (int i = 0; i < historyUUIDs.length; i++) {
                if (historyUUIDs[i] == null) {
                    onFailure.accept(new IllegalArgumentException("no history uuid for [" + leaderIndex + "][" + i + "]"));
                    return;
                }
            }
            historyUUIDConsumer.accept(historyUUIDs);
        };
        IndicesStatsRequest request = new IndicesStatsRequest();
        request.clear();
        request.indices(leaderIndex);
        leaderClient.admin().indices().stats(request, ActionListener.wrap(indicesStatsHandler, onFailure));
    }

    public static Client wrapClient(Client client, Map<String, String> headers) {
        if (headers.isEmpty()) {
            return client;
        } else {
            final ThreadContext threadContext = client.threadPool().getThreadContext();
            Map<String, String> filteredHeaders = headers.entrySet().stream()
                .filter(e -> ShardFollowTask.HEADER_FILTERS.contains(e.getKey()))
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
            return new FilterClient(client) {
                @Override
                protected <Request extends ActionRequest, Response extends ActionResponse>
                void doExecute(Action<Response> action, Request request, ActionListener<Response> listener) {
                    final Supplier<ThreadContext.StoredContext> supplier = threadContext.newRestorableContext(false);
                    try (ThreadContext.StoredContext ignore = stashWithHeaders(threadContext, filteredHeaders)) {
                        super.doExecute(action, request, new ContextPreservingActionListener<>(supplier, listener));
                    }
                }
            };
        }
    }

    private static ThreadContext.StoredContext stashWithHeaders(ThreadContext threadContext, Map<String, String> headers) {
        final ThreadContext.StoredContext storedContext = threadContext.stashContext();
        threadContext.copyHeaders(headers.entrySet());
        return storedContext;
    }

    private static ElasticsearchStatusException indexMetadataNonCompliantRemoteLicense(
            final String leaderIndex, final RemoteClusterLicenseChecker.LicenseCheck licenseCheck) {
        final String clusterAlias = licenseCheck.remoteClusterLicenseInfo().clusterAlias();
        final String message = String.format(
                Locale.ROOT,
                "can not fetch remote index [%s:%s] metadata as the remote cluster [%s] is not licensed for [ccr]; %s",
                clusterAlias,
                leaderIndex,
                clusterAlias,
                RemoteClusterLicenseChecker.buildErrorMessage(
                        "ccr",
                        licenseCheck.remoteClusterLicenseInfo(),
                        RemoteClusterLicenseChecker::isLicensePlatinumOrTrial));
        return new ElasticsearchStatusException(message, RestStatus.BAD_REQUEST);
    }

    private static ElasticsearchStatusException clusterStateNonCompliantRemoteLicense(
            final RemoteClusterLicenseChecker.LicenseCheck licenseCheck) {
        final String clusterAlias = licenseCheck.remoteClusterLicenseInfo().clusterAlias();
        final String message = String.format(
                Locale.ROOT,
                "can not fetch remote cluster state as the remote cluster [%s] is not licensed for [ccr]; %s",
                clusterAlias,
                RemoteClusterLicenseChecker.buildErrorMessage(
                        "ccr",
                        licenseCheck.remoteClusterLicenseInfo(),
                        RemoteClusterLicenseChecker::isLicensePlatinumOrTrial));
        return new ElasticsearchStatusException(message, RestStatus.BAD_REQUEST);
    }

    private static ElasticsearchStatusException indexMetadataUnknownRemoteLicense(
            final String leaderIndex, final String clusterAlias, final Exception cause) {
        final String message = String.format(
                Locale.ROOT,
                "can not fetch remote index [%s:%s] metadata as the license state of the remote cluster [%s] could not be determined",
                clusterAlias,
                leaderIndex,
                clusterAlias);
        return new ElasticsearchStatusException(message, RestStatus.BAD_REQUEST, cause);
    }

    private static ElasticsearchStatusException clusterStateUnknownRemoteLicense(final String clusterAlias, final Exception cause) {
        final String message = String.format(
                Locale.ROOT,
                "can not fetch remote cluster state as the license state of the remote cluster [%s] could not be determined", clusterAlias);
        return new ElasticsearchStatusException(message, RestStatus.BAD_REQUEST, cause);
    }

}
