/*
 * Copyright 2018-2021 Bloomberg Finance LP
 *
 * Licensed 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.
 */

#ifndef INCLUDED_BUILDBOXCOMMON_GRPCCLIENT
#define INCLUDED_BUILDBOXCOMMON_GRPCCLIENT

#include <functional>
#include <memory>

#include <buildboxcommon_connectionoptions.h>
#include <buildboxcommon_grpcerror.h>
#include <buildboxcommon_grpcretrier.h>
#include <buildboxcommon_protos.h>
#include <buildboxcommon_requestmetadata.h>
#include <buildboxcommon_schedulingmetadata.h>

namespace buildboxcommon {

/**
 * Implements a mechanism to communicate with gRPC servers.
 */
class GrpcClient {
  public:
    // Helper method to shutdown and clear out a CompletionQueue must be
    // called when going out of scope to avoid a memory leak.
    static void ShutdownAndEmptyQueue(grpc::CompletionQueue *cq)
    {
        cq->Shutdown();
        // Drain queue
        void *ignored_tag = nullptr;
        bool ignored_ok = false;
        while (cq->Next(&ignored_tag, &ignored_ok)) {
        }
    }

    GrpcClient() {}

    /**
     * Connect to the gRPC endpoint with the given connection options.
     */
    void init(const ConnectionOptions &options);

    std::shared_ptr<grpc::Channel> channel() { return d_channel; }

    void setToolDetails(const std::string &tool_name,
                        const std::string &tool_version);
    /**
     * Set the optional ID values to be attached to requests.
     */
    void setRequestMetadata(const std::string &action_id,
                            const std::string &tool_invocation_id,
                            const std::string &correlated_invocations_id,
                            const std::string &action_mnemonic = "",
                            const std::string &target_id = "",
                            const std::string &configuration_id = "");
    void setActionId(const std::string &action_id);
    void setSchedulingMetadata(
        build::buildgrid::SchedulingMetadata scheduling_metadata);

    /**
     * Set the function used to attach metadata to requests. If this is called
     * to change from the default behaviour, the any client-level metadata set
     * with `setRequestMetadata` will be ignored.
     */
    void setMetadataAttacher(std::function<void(grpc::ClientContext *)> fn);

    /**
     * Set the function used to get the server context of the current request.
     * This only applies to gRPC servers that send requests to other gRPC
     * servers. If set, the server context returned by this function is used as
     * basis for constructing client contexts, enabling deadline and
     * cancellation propagation.
     */
    void
    setServerContextGetter(std::function<const grpc::ServerContext *()> fn);

    /* Passing this optional object to methods allows callers to read stats
     * about the gRPC request.
     */
    struct RequestStats {
        RequestStats() = default;

        // Number of gRPC retries (requests after the initial one) issued
        // until the call succeeded, failed with a non-retryable error code, or
        // the retry limit was exceeded.
        unsigned int d_grpcRetryCount = 0;
    };

    std::string instanceName() const;

    void setInstanceName(const std::string &instance_name);

    void
    issueRequest(const buildboxcommon::GrpcRetrier::GrpcInvocation &invocation,
                 const std::string &invocationName,
                 RequestStats *requestStats) const;

    void
    issueRequest(const buildboxcommon::GrpcRetrier::GrpcInvocation &invocation,
                 const std::string &invocationName,
                 const std::chrono::seconds &requestTimeout,
                 RequestStats *requestStats) const;

    void
    issueRequest(const buildboxcommon::GrpcRetrier::GrpcInvocation &invocation,
                 const std::string &invocationName, int64_t sizeBytes,
                 RequestStats *requestStats) const;

    GrpcRetrier
    makeRetrier(const GrpcRetrier::GrpcInvocation &invocation,
                const std::string &name,
                const std::chrono::seconds &requestTimeout =
                    std::chrono::seconds::zero(),
                const std::set<grpc::StatusCode> &retryOnCodes = {}) const;

    GrpcRetrier
    makeRetrier(const GrpcRetrier::GrpcInvocation &invocation,
                const std::string &name, int64_t sizeBytes,
                const std::set<grpc::StatusCode> &retryOnCodes = {}) const;

    int retryLimit() const;

    std::chrono::seconds requestTimeout() const;

    void setRetryLimit(int limit);

    void setRequestTimeout(const std::chrono::seconds &requestTimeout);

    static size_t maxMessageSizeBytes();

  private:
    // initialized here to prevent errors, in case options are not passed into
    // init
    int d_grpcRetryLimit = 0;
    int d_grpcRetryDelay =
        100; // NOLINT (cppcoreguidelines-avoid-magic-numbers)

    std::set<grpc::StatusCode> d_retryOnCodes;

    std::chrono::seconds d_grpcRequestTimeout = std::chrono::seconds::zero();
    int64_t d_grpcMinThroughput = 0;

    std::shared_ptr<grpc::Channel> d_channel;
    std::string d_instanceName;

    RequestMetadataGenerator d_metadata_generator;
    SchedulingMetadataGenerator d_schedulingMetadataGenerator;
    std::function<void(grpc::ClientContext *)> d_metadata_attach_function =
        [&](grpc::ClientContext *context) {
            d_metadata_generator.attach_request_metadata(context);
            d_schedulingMetadataGenerator.attachSchedulingMetadata(context);
        };

    std::function<const grpc::ServerContext *()> d_serverContextGetter = []() {
        return nullptr;
    };

    // GRPC constant which is a server-side receive value or a client side send
    // value, minus an arbitrary delta for GRPC metadata
    // https://github.com/grpc/grpc/blob/master/include/grpc/impl/grpc_types.h
    static const size_t s_maxMessageSizeBytes =
        GRPC_DEFAULT_MAX_RECV_MESSAGE_LENGTH - (1 << 16);
};

} // namespace buildboxcommon

#endif
