// Copyright 2020 Bloomberg Finance L.P
//
// 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_GRPCRETRIER
#define INCLUDED_BUILDBOXCOMMON_GRPCRETRIER

#include <buildboxcommon_protos.h>
#include <buildboxcommon_requestmetadata.h>

#include <buildboxcommon_connectionoptions.h>

#include <chrono>
#include <functional>
#include <set>
#include <string>

namespace buildboxcommon {

class GrpcRetrier final {

    /* This class wraps a function that issues a gRPC request and will attempt
     * to invoke the request until it either succeeds, fails with a final,
     * non-retryable error, or the limit of attempts was exceeded.
     *
     * Between retry attempts it will use an exponential-backoff delay.
     *
     * Sample usage:
     *
     * ```
     * const unsigned int retryLimit = 3;
     * const std::chrono::milliseconds retryDelayBase(100);
     *
     * GrpcRetrier r(retryLimit, retryDelayBase, grpcInvocationFunction,
     *               "grpcInvocationName()");
     *
     * if (r.issueRequest()) {
     *    // Received a final answer before exceeding the retry limit.
     *
     *    if (r.status().ok()) {
     *        // The gRPC request was successful.
     *    }
     *    else {
     *        // The request failed with a final error, such as `NOT_FOUND`.
     *    }
     *
     * }
     * else {
     *    // Retry limit exceeded; gRPC request failed.
     * }
     * ```
     *
     */

  public:
    typedef std::function<grpc::Status(grpc::ClientContext &)> GrpcInvocation;

    typedef std::function<void(grpc::ClientContext *)> MetadataAttacher;

    typedef std::set<grpc::StatusCode> GrpcStatusCodes;

    typedef std::set<std::pair<grpc::StatusCode, std::string>>
        GrpcStatusCodeMessagePairs;

    GrpcRetrier(const unsigned int retryLimit,
                const std::chrono::milliseconds &retryDelayBase,
                const GrpcInvocation &grpcInvocation,
                const std::string &grpcInvocationName)
        : GrpcRetrier(retryLimit, retryDelayBase, grpcInvocation,
                      grpcInvocationName, {}, nullptr)
    {
    }

    GrpcRetrier(const unsigned int retryLimit,
                const std::chrono::milliseconds &retryDelayBase,
                const GrpcInvocation &grpcInvocation,
                const std::string &grpcInvocationName,
                const GrpcStatusCodes &retryableStatusCodes,
                const MetadataAttacher &metadataAttacher = nullptr,
                const std::chrono::seconds &requestTimeout =
                    std::chrono::seconds::zero())
        : GrpcRetrier(retryLimit, retryDelayBase, grpcInvocation,
                      grpcInvocationName, retryableStatusCodes, {},
                      metadataAttacher, requestTimeout)
    {
    }

    GrpcRetrier(
        const unsigned int retryLimit,
        const std::chrono::milliseconds &retryDelayBase,
        const GrpcInvocation &grpcInvocation,
        const std::string &grpcInvocationName,
        const GrpcStatusCodes &retryableStatusCodes,
        const GrpcStatusCodeMessagePairs &retryableStatusCodeMessagePairs,
        const MetadataAttacher &metadataAttacher = nullptr,
        const std::chrono::seconds &requestTimeout =
            std::chrono::seconds::zero(),
        const grpc::ServerContext *serverContext = nullptr)
        : d_grpcInvocation(grpcInvocation),
          d_grpcInvocationName(grpcInvocationName), d_retryLimit(retryLimit),
          d_retryDelayBase(retryDelayBase),
          d_retryableStatusCodes(retryableStatusCodes),
          d_retryableStatusCodeMessagePairs(retryableStatusCodeMessagePairs),
          d_metadataAttacher(metadataAttacher), d_serverContext(serverContext),
          d_retryAttempts(0), d_requestTimeout(requestTimeout)
    {

        // Always retry on UNAVAILABLE
        d_retryableStatusCodes.insert(grpc::StatusCode::UNAVAILABLE);

        // Always retry on stream removed
        d_retryableStatusCodeMessagePairs.insert(
            {grpc::StatusCode::UNKNOWN, "Stream removed"});
        // OK status is always OK
        d_okStatusCodes.insert(grpc::StatusCode::OK);
    }

    // default destructor
    ~GrpcRetrier() = default;

    // delete move and copy constructors
    GrpcRetrier(const GrpcRetrier &) = delete;
    GrpcRetrier &operator=(const GrpcRetrier &) = delete;
    GrpcRetrier(GrpcRetrier &&) = delete;
    GrpcRetrier &operator=(GrpcRetrier &&) = delete;

    /* Returns the maximum number of retries that will be attempted after
     * an initial request that fails.
     */
    const unsigned int &retryLimit() const { return d_retryLimit; }

    /* Returns the value used as a base for the exponential-backoff wait
     * between attempts.
     */
    const std::chrono::milliseconds &retryDelayBase() const
    {
        return d_retryDelayBase;
    }

    // Returns the request timeout. 0 indicates no timeout.
    const std::chrono::seconds &requestTimeout() const
    {
        return d_requestTimeout;
    }

    /* Issue the gRPC request and return whether the request was completed in
     * fewer retries than the limit. (Note that the request might have failed
     * with a non-retryable status.)
     *
     * If the retry count was exceeded, this method returns `false`.
     *
     * `status()` will return the last value returned by the gRPC request.
     */
    bool issueRequest();

    /* Set of codes that enable to retry the request.  (Should contain errors
     * that are transient.)
     */
    const GrpcStatusCodes &retryableStatusCodes() const
    {
        return d_retryableStatusCodes;
    }

    /* Add an additional grpc::StatusCode to the set of codes to retry */
    void addRetryableStatusCode(const grpc::StatusCode retryableCode)
    {
        d_retryableStatusCodes.insert(retryableCode);
    }

    /* Set of pairs of codes and messages that enable to retry the request.
     * (Should contain errors that are transient.)
     */
    const GrpcStatusCodeMessagePairs &retryableStatusCodeMessagePairs() const
    {
        return d_retryableStatusCodeMessagePairs;
    }

    /* Add a pair of {grpc::StatusCode, std::string} to the set of codes to
     * retry */
    void addRetryableStatusCodeMessagePair(
        const std::pair<grpc::StatusCode, std::string>
            retryableStatusCodeMessagePair)
    {
        d_retryableStatusCodeMessagePairs.insert(
            retryableStatusCodeMessagePair);
    }

    /* Set of codes that are "OK" from the retriers perspective. An example
     * might be a client treating ALREADY_EXISTS as an OK status when trying to
     * upload a blob.
     */
    const GrpcStatusCodes &okStatusCodes() const { return d_okStatusCodes; }

    /* Add an additional grpc::StatusCode to the set of "OK" codes */
    void addOkStatusCode(const grpc::StatusCode okCode)
    {
        d_okStatusCodes.insert(okCode);
    }

    void setMetadataAttacher(const MetadataAttacher &metadataAttacher)
    {
        d_metadataAttacher = metadataAttacher;
    }

    /* Return the `grpc::Status` received on the last attempt. */
    grpc::Status status() const { return d_status; }

    /* Number of retries attempted in `issueRequest()` */
    const unsigned int &retryAttempts() const { return d_retryAttempts; }

  private:
    // gRPC callback to perform a request and its human-readable name for logs:

    const GrpcInvocation d_grpcInvocation;
    const std::string d_grpcInvocationName;

    // Maximum number of attempts and delay between retries:
    const unsigned int d_retryLimit;

    std::chrono::milliseconds d_retryDelayBase;
    // The value set in the constructor for this variable will be used unless
    // the server attaches a `RetryInfo` with a non-zero `Duration in the first
    // error Status, in which case it will be overwritten.
    //
    // According to `google.rpc.RetryInfo` docs:
    // "Clients should wait until `retry_delay` [...] has passed since
    // receiving the error response before retrying.
    // If retrying requests also fail, clients should use an exponential
    // backoff [...] based on  `retry_delay`, until either a maximum number of
    // retries have been reached or a maximum retry delay cap has been
    // reached."

    // Status codes to retry:
    GrpcStatusCodes d_retryableStatusCodes;

    // Pairs of status codes & messages to retry:
    GrpcStatusCodeMessagePairs d_retryableStatusCodeMessagePairs;

    // Status codes to consider successes
    GrpcStatusCodes d_okStatusCodes;

    // Optional callback to attach metadata to the request before issuing it:
    MetadataAttacher d_metadataAttacher;

    // Optional server context to use for constructing client contexts.
    const grpc::ServerContext *d_serverContext = nullptr;

    // Results after `issueRequest()`:
    grpc::Status d_status;        // Last status received from the server.
    unsigned int d_retryAttempts; // Number of retries performed (excluding
                                  // original request):

    // How long to wait for a response per gRPC request invocation
    // 0 indicates no timeout.
    std::chrono::seconds d_requestTimeout;

    // Returns whether the status can be retried or should be considered final.
    bool statusIsRetryable(const grpc::Status &status) const;

    // Returns whether the status is an "ok" status from the retrier's
    // perspective.
    bool statusIsOk(const grpc::Status &status) const;
};

} // namespace buildboxcommon

#endif
