// 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.

#include <buildboxcommon_grpcretrier.h>

#include <buildboxcommon_logging.h>
#include <buildboxcommon_stringutils.h>

#include <google/protobuf/util/time_util.h>
#include <google/rpc/error_details.grpc.pb.h>

#include <cmath>
#include <grpcpp/support/status.h>
#include <thread>

namespace {
std::string retryingInvocationWarningMessage(
    const std::string &grpcInvocationName, const grpc::Status &grpcError,
    const unsigned int attemptNumber, const unsigned int totalAttempts,
    const double &retryDelayMs)
{
    std::stringstream s;
    s << "Attempt " << attemptNumber + 1 << "/" << totalAttempts + 1;

    if (!grpcInvocationName.empty()) {
        s << " for \"" << grpcInvocationName << "\"";
    }

    s << " failed with gRPC error [" << grpcError.error_code() << ": "
      << grpcError.error_message() << "], "
      << "retrying in " << retryDelayMs << " ms...";

    return s.str();
}

std::string
retryAttemptsExceededErrorMessage(const std::string &grpcInvocationName,
                                  const grpc::Status &grpcError,
                                  const unsigned int retryLimit)
{
    std::stringstream s;
    s << "Retry limit (" << retryLimit << ") exceeded";
    if (!grpcInvocationName.empty()) {
        s << " for \"" << grpcInvocationName << "\"";
    }

    s << ", last gRPC error was [" << grpcError.error_code() << ": "
      << grpcError.error_message() << "]";

    return s.str();
}

} // namespace

namespace buildboxcommon {
bool GrpcRetrier::issueRequest()
{
    d_retryAttempts = 0;

    // Apply the request timeout across all attempts for a given RPC as per
    // https://github.com/grpc/proposal/blob/master/A6-client-retries.md
    std::chrono::system_clock::time_point deadline;
    if (d_requestTimeout > std::chrono::seconds::zero()) {
        deadline = std::chrono::system_clock::now() + d_requestTimeout;
    }

    while (true) {
        std::unique_ptr<grpc::ClientContext> context;
        if (d_serverContext) {
            if (d_serverContext->IsCancelled()) {
                d_status = grpc::Status::CANCELLED;
                return true;
            }
            context = grpc::ClientContext::FromServerContext(*d_serverContext);
        }
        else {
            context = std::make_unique<grpc::ClientContext>();
        }

        if (d_metadataAttacher) {
            d_metadataAttacher(context.get());
        }

        if (d_requestTimeout > std::chrono::seconds::zero()) {
            context->set_deadline(deadline);
        }

        d_status = d_grpcInvocation(*context);

        if (statusIsOk(d_status)) {
            if (d_retryAttempts > 0) {
                BUILDBOX_LOG_INFO(
                    "\"" << d_grpcInvocationName << "\" returned okay status "
                         << std::to_string(d_status.error_code()) << " on "
                         << StringUtils::ordinal(d_retryAttempts + 1)
                         << " attempt");
            }

            return true;
        }

        if (!statusIsRetryable(d_status)) {
            std::string extraLogContext = "";
            if (d_status.error_code() == grpc::DEADLINE_EXCEEDED &&
                d_requestTimeout > std::chrono::seconds::zero()) {
                bool server_timeout =
                    (std::chrono::system_clock::now() < deadline);
                if (server_timeout) {
                    extraLogContext = " (server timeout)";
                }
                else {
                    extraLogContext = " (client timeout)";
                }
            }

            if (d_status.error_code() == grpc::CANCELLED) {
                BUILDBOX_LOG_WARNING(
                    d_grpcInvocationName +
                    " cancelled: " + std::to_string(d_status.error_code()) +
                    ": " + d_status.error_message() + extraLogContext);
            }
            else {
                BUILDBOX_LOG_ERROR(
                    d_grpcInvocationName +
                    " failed with: " + std::to_string(d_status.error_code()) +
                    ": " + d_status.error_message() + extraLogContext);
            }
            return true;
        }

        // The error might contain a `RetryInfo` message specifying a number of
        // seconds to wait before retrying. If so, use it for the base value.
        if (d_retryAttempts == 0 && !d_status.error_details().empty()) {
            // grpc::Status.error_details() is expected to be a serialized
            // google.rpc.Status message
            // https://github.com/grpc/grpc/blob/39d007e83bf5f2eebd0a27fd0bbd71435639dbeb/include/grpcpp/impl/codegen/status.h#L114
            google::rpc::Status error_details;
            if (error_details.ParseFromString(d_status.error_details())) {
                google::rpc::RetryInfo retryInfo;
                // This is a repeated Any field, so go through each one
                // and see if there's a RetryInfo in there
                for (const google::protobuf::Any &detailed_error :
                     error_details.details()) {
                    if (detailed_error.UnpackTo(&retryInfo)) {
                        const google::protobuf::int64 serverDelay = google::
                            protobuf::util::TimeUtil::DurationToMilliseconds(
                                retryInfo.retry_delay());

                        if (serverDelay > 0) {
                            d_retryDelayBase =
                                std::chrono::milliseconds(serverDelay);
                            BUILDBOX_LOG_DEBUG(
                                "Overriding retry delay base with "
                                "value specified by server: "
                                << d_retryDelayBase.count() << " ms");
                        }
                    }
                }
            }
        }

        // The call failed and could be retryable on its own.
        if (d_retryAttempts >= d_retryLimit) {
            const std::string errorMessage = retryAttemptsExceededErrorMessage(
                d_grpcInvocationName, d_status, d_retryLimit);
            BUILDBOX_LOG_ERROR(errorMessage);

            return false;
        }

        // Delay the next call based on the number of attempts made:
        const auto retryDelay = d_retryDelayBase * pow(1.6, d_retryAttempts);
        BUILDBOX_LOG_WARNING(retryingInvocationWarningMessage(
            d_grpcInvocationName, d_status, d_retryAttempts, d_retryLimit,
            retryDelay.count()));

        const auto retryTime = std::chrono::system_clock::now() + retryDelay;

        // If the retry delay will exceed the deadline, abort immediately.
        if (d_requestTimeout > std::chrono::seconds::zero() &&
            retryTime >= deadline) {
            BUILDBOX_LOG_ERROR(d_grpcInvocationName + " failed with: " +
                               std::to_string(d_status.error_code()) + ": " +
                               d_status.error_message() + " (client timeout)");
            return true;
        }

        if (d_serverContext) {
            // Periodically check whether the request was cancelled.
            while (true) {
                if (d_serverContext->IsCancelled()) {
                    d_status = grpc::Status::CANCELLED;
                    return true;
                }
                const auto nextCancelCheck = std::chrono::system_clock::now() +
                                             std::chrono::milliseconds(500);
                if (retryTime <= nextCancelCheck) {
                    break;
                }
                std::this_thread::sleep_until(nextCancelCheck);
            }
        }

        std::this_thread::sleep_until(retryTime);

        d_retryAttempts++;
    }
}

bool GrpcRetrier::statusIsRetryable(const grpc::Status &status) const
{
    // Check if the status code is in the retryable set of status codes
    // or if in the message pair the error code is the same and the error
    // message starts with the error string in message pair
    return d_retryableStatusCodes.count(status.error_code()) ||
           std::any_of(d_retryableStatusCodeMessagePairs.begin(),
                       d_retryableStatusCodeMessagePairs.end(),
                       [&status](const auto &pair) {
                           return status.error_code() == pair.first &&
                                  status.error_message().compare(
                                      0, pair.second.size(), pair.second) == 0;
                       });
}

bool GrpcRetrier::statusIsOk(const grpc::Status &status) const
{
    return d_okStatusCodes.count(status.error_code());
}

} // namespace buildboxcommon
