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

#include <buildboxcommon_exception.h>
#include <buildboxcommon_grpcclient.h>
#include <buildboxcommon_logging.h>

#include <grpc/grpc.h>
#include <string>
#include <utility>

namespace buildboxcommon {

void GrpcClient::init(const ConnectionOptions &options)
{
    std::shared_ptr<grpc::Channel> channel = options.createChannel();
    this->d_grpcRetryLimit = std::stoi(options.d_retryLimit);
    this->d_grpcRetryDelay = std::stoi(options.d_retryDelay);
    this->d_grpcRequestTimeout =
        std::chrono::seconds(std::stoi(options.d_requestTimeout));
    this->d_retryOnCodes = options.d_retryOnCodes;
    this->d_grpcMinThroughput = options.d_minThroughput;
    this->d_channel = channel;
    this->d_instanceName = options.d_instanceName;
}

std::string GrpcClient::instanceName() const { return d_instanceName; }

void GrpcClient::setInstanceName(const std::string &instance_name)
{
    d_instanceName = instance_name;
}

void GrpcClient::setToolDetails(const std::string &tool_name,
                                const std::string &tool_version)
{
    d_metadata_generator.set_tool_details(tool_name, tool_version);
}

void GrpcClient::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)
{
    d_metadata_generator.set_action_id(action_id);
    d_metadata_generator.set_tool_invocation_id(tool_invocation_id);
    d_metadata_generator.set_correlated_invocations_id(
        correlated_invocations_id);
    d_metadata_generator.set_action_mnemonic(action_mnemonic);
    d_metadata_generator.set_target_id(target_id);
    d_metadata_generator.set_configuration_id(configuration_id);
}

void GrpcClient::setActionId(const std::string &action_id)
{
    d_metadata_generator.set_action_id(action_id);
}

void GrpcClient::setSchedulingMetadata(
    build::buildgrid::SchedulingMetadata scheduling_metadata)
{
    d_schedulingMetadataGenerator.setSchedulingMetadata(
        std::move(scheduling_metadata));
}

void GrpcClient::issueRequest(const GrpcRetrier::GrpcInvocation &invocation,
                              const std::string &invocationName,
                              RequestStats *requestStats) const
{
    issueRequest(invocation, invocationName, std::chrono::seconds::zero(),
                 requestStats);
}

void GrpcClient::issueRequest(const GrpcRetrier::GrpcInvocation &invocation,
                              const std::string &invocationName,
                              const std::chrono::seconds &requestTimeout,
                              RequestStats *requestStats) const
{
    auto retrier = makeRetrier(invocation, invocationName, requestTimeout,
                               d_retryOnCodes);
    retrier.issueRequest();

    if (requestStats != nullptr) {
        requestStats->d_grpcRetryCount += retrier.retryAttempts();
        // (Adding instead of assigning so that a running total can be kept in
        // a single struct.)
    }

    if (!retrier.status().ok()) {
        GrpcError::throwGrpcError(retrier.status());
    }
}

void GrpcClient::issueRequest(const GrpcRetrier::GrpcInvocation &invocation,
                              const std::string &invocationName,
                              int64_t sizeBytes,
                              RequestStats *requestStats) const
{
    auto retrier =
        makeRetrier(invocation, invocationName, sizeBytes, d_retryOnCodes);
    retrier.issueRequest();

    if (requestStats != nullptr) {
        requestStats->d_grpcRetryCount += retrier.retryAttempts();
        // (Adding instead of assigning so that a running total can be kept in
        // a single struct.)
    }

    if (!retrier.status().ok()) {
        GrpcError::throwGrpcError(retrier.status());
    }
}

GrpcRetrier
GrpcClient::makeRetrier(const GrpcRetrier::GrpcInvocation &invocation,
                        const std::string &invocationName,
                        const std::chrono::seconds &requestTimeout,
                        const std::set<grpc::StatusCode> &retryOnCodes) const
{
    // Pick the minimum non-zero timeout (from connectionoptions or override)
    auto min_nonzero_comp = [](std::chrono::seconds a,
                               std::chrono::seconds b) {
        bool a_nonzero = a != std::chrono::seconds::zero();
        bool b_nonzero = b != std::chrono::seconds::zero();
        return (a_nonzero && b_nonzero && a < b) || (!b_nonzero);
    };

    const std::chrono::seconds &shortestRequestTimeout =
        std::min(requestTimeout, d_grpcRequestTimeout, min_nonzero_comp);

    // Guaranteed RVO which does not need copy or move constructors
    return GrpcRetrier(d_grpcRetryLimit,
                       std::chrono::milliseconds(d_grpcRetryDelay), invocation,
                       invocationName, retryOnCodes, {},
                       d_metadata_attach_function, shortestRequestTimeout,
                       d_serverContextGetter());
}

GrpcRetrier
GrpcClient::makeRetrier(const GrpcRetrier::GrpcInvocation &invocation,
                        const std::string &invocationName, int64_t sizeBytes,
                        const std::set<grpc::StatusCode> &retryOnCodes) const
{
    std::chrono::seconds timeout = d_grpcRequestTimeout;
    if (timeout > std::chrono::seconds::zero() && sizeBytes > 0) {
        if (d_grpcMinThroughput > 0) {
            // Round up to make sure the timeout doesn't trigger too early
            timeout += std::chrono::seconds(
                (sizeBytes + d_grpcMinThroughput - 1) / d_grpcMinThroughput);
        }
        else {
            // Without a configured minimum throughput, calculate the number
            // of messages required to transfer the specified number of bytes,
            // and then grant each message as much time as a non-streaming
            // call.
            const uint64_t nMessages = (static_cast<uint64_t>(sizeBytes) +
                                        s_maxMessageSizeBytes - 1) /
                                       s_maxMessageSizeBytes;
            timeout = d_grpcRequestTimeout * nMessages;
        }
    }

    // Guaranteed RVO which does not need copy or move constructors
    return GrpcRetrier(
        d_grpcRetryLimit, std::chrono::milliseconds(d_grpcRetryDelay),
        invocation, invocationName, retryOnCodes, {},
        d_metadata_attach_function, timeout, d_serverContextGetter());
}

int GrpcClient::retryLimit() const { return d_grpcRetryLimit; }

std::chrono::seconds GrpcClient::requestTimeout() const
{
    return d_grpcRequestTimeout;
}

void GrpcClient::setRetryLimit(int limit) { d_grpcRetryLimit = limit; }

void GrpcClient::setRequestTimeout(const std::chrono::seconds &requestTimeout)
{
    d_grpcRequestTimeout = requestTimeout;
}

void GrpcClient::setMetadataAttacher(
    std::function<void(grpc::ClientContext *)> fn)
{
    d_metadata_attach_function = std::move(fn);
}

void GrpcClient::setServerContextGetter(
    std::function<const grpc::ServerContext *()> fn)
{
    d_serverContextGetter = std::move(fn);
}

size_t GrpcClient::maxMessageSizeBytes() { return s_maxMessageSizeBytes; }

} // namespace buildboxcommon
