/*
 * Copyright 2020 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_BUILDBOXWORKER_BOTSESSIONUTILS
#define INCLUDED_BUILDBOXWORKER_BOTSESSIONUTILS

#include <buildboxcommon_connectionoptions.h>
#include <buildboxcommon_grpcretrier.h>
#include <buildboxcommon_protos.h>
#include <buildboxcommon_requestmetadata.h>
#include <buildboxworker_worker.h>

#include <build/bazel/remote/execution/v2/remote_execution.grpc.pb.h>
#include <google/devtools/remoteworkers/v1test2/bots.grpc.pb.h>

namespace buildboxworker {

namespace proto {
using namespace google::devtools::remoteworkers::v1test2;
} // namespace proto

struct BotSessionUtils {
    // Metadata types:
    typedef std::multimap<grpc::string_ref, grpc::string_ref>
        GrpcServerMetadataMultiMap;

    typedef std::unordered_multimap<buildboxcommon::Digest,
                                    buildboxcommon::ExecuteOperationMetadata>
        ExecuteOperationMetadataEntriesMultiMap;

    /**
     * Helper that calls the UpdateBotSession grpc method.
     *
     * This retries the UpdateBotSession call on failures a set number of
     * times and exits if it still failed afterwards. The exiting can be
     * disabled by setting should_exit to false.
     *
     */
    static grpc::Status updateBotSession(
        const proto::BotStatus botStatus,
        std::shared_ptr<proto::Bots::StubInterface> stub,
        const buildboxcommon::ConnectionOptions &botsServerConnection,
        proto::BotSession *session,
        const buildboxcommon::GrpcRetrier::GrpcStatusCodes &statusCodes,
        const buildboxcommon::RequestMetadataGenerator &requestMetadata,
        const RunnerPartialExecutionMetadata &partialExecutionMetadata,
        std::unique_lock<std::mutex> *sessionLock,
        ExecuteOperationMetadataEntriesMultiMap
            *executeOperationMetadataEntries = nullptr);

    static grpc::Status updateBotSession(
        const proto::BotStatus botStatus,
        std::shared_ptr<proto::Bots::StubInterface> stub,
        const buildboxcommon::ConnectionOptions &botsServerConnection,
        proto::BotSession *session,
        const buildboxcommon::GrpcRetrier::GrpcStatusCodes &statusCodes,
        const std::atomic_bool &cancelCall,
        const buildboxcommon::RequestMetadataGenerator &requestMetadata,
        const RunnerPartialExecutionMetadata &partialExecutionMetadata,
        std::unique_lock<std::mutex> *sessionLock,
        ExecuteOperationMetadataEntriesMultiMap
            *executeOperationMetadataEntries = nullptr);

    static grpc::Status createBotSession(
        std::shared_ptr<proto::Bots::StubInterface> stub,
        const buildboxcommon::ConnectionOptions &botsServerConnection,
        proto::BotSession *session,
        const buildboxcommon::GrpcRetrier::GrpcStatusCodes &statusCodes,
        const buildboxcommon::RequestMetadataGenerator &requestMetadata);

    static grpc::Status createBotSession(
        std::shared_ptr<proto::Bots::StubInterface> stub,
        const buildboxcommon::ConnectionOptions &botsServerConnection,
        proto::BotSession *session,
        const buildboxcommon::GrpcRetrier::GrpcStatusCodes &statusCodes,
        const buildboxcommon::RequestMetadataGenerator &requestMetadata,
        const std::atomic_bool &cancelCall);

    /**
     * Given a multimap obtained by calling
     * `grpc::ClientContext.GetServerTrailingMetadata()`, searches for metadata
     * entries under the "executeoperationmetadata-bin" name and returns a list
     * with the `ExecuteOperationMetadata` messages they contain.
     */
    static ExecuteOperationMetadataEntriesMultiMap
    extractExecuteOperationMetadata(
        const GrpcServerMetadataMultiMap &metadata);

    /**
     * Given a CompletionQueue waiting for a final response
     * wait until the final message has been recieved or
     * cancelCall is set, in which case the call is cancelled
     * and the CompletionQueue shutdown.
     */
    static void waitForResponseOrCancel(grpc::ClientContext *context,
                                        grpc::CompletionQueue *cq,
                                        grpc::Status *status,
                                        const std::atomic_bool &cancelCall);
};

} // namespace buildboxworker

#endif
