/*
 * Copyright 2018 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 <buildboxworker_botsessionutils.h>
#include <buildboxworker_metricnames.h>
#include <buildboxworker_worker.h>

#include <buildboxcommon_grpcclient.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_requestmetadata.h>
#include <buildboxcommon_timeutils.h>
#include <buildboxcommonmetrics_durationmetrictimer.h>
#include <buildboxcommonmetrics_metricguard.h>

#include <utility>

namespace buildboxworker {

using namespace buildboxcommon;

grpc::Status BotSessionUtils::updateBotSession(
    const proto::BotStatus botStatus,
    std::shared_ptr<proto::Bots::StubInterface> stub,
    const buildboxcommon::ConnectionOptions &botsServerConnection,
    proto::BotSession *session,
    const GrpcRetrier::GrpcStatusCodes &statusCodes,
    const buildboxcommon::RequestMetadataGenerator &requestMetadata,
    const RunnerPartialExecutionMetadata &partialExecutionMetadata,
    std::unique_lock<std::mutex> *sessionLock,
    ExecuteOperationMetadataEntriesMultiMap *executeOperationMetadataEntries)
{
    std::atomic_bool shouldCancel(false);
    return updateBotSession(botStatus, std::move(stub), botsServerConnection,
                            session, statusCodes, shouldCancel,
                            requestMetadata, partialExecutionMetadata,
                            sessionLock, executeOperationMetadataEntries);
}

std::string sessionDebugInfo(proto::BotSession *session)
{
    std::string info;
    info += " bot_name=[" + session->name() + "]";
    info += " bot_id=[" + session->bot_id() + "]";
    info += " bot_status=[" + BotStatus_Name(session->status()) + "]";
    for (auto &lease : session->leases()) {
        info += " lease_id=[" + lease.id() + "]";
        info += " lease_state=[" + LeaseState_Name(lease.state()) + "]";
    }
    return info;
}

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

    updateRequest.set_name(tempSession.name());
    *updateRequest.mutable_bot_session() = tempSession;
    updateRequest.mutable_bot_session()->set_status(botStatus);

    const auto updateInvocation = [&](grpc::ClientContext &context) {
        // If cancelCall is already set, return cancelled right away
        // instead of sending the request.
        if (cancelCall) {
            return grpc::Status(
                grpc::CANCELLED,
                "Cancellation requested before sending request");
        }

        BUILDBOX_LOG_DEBUG("Updating bot session."
                           << sessionDebugInfo(&tempSession));

        // Unlock the session lock whilst we make the UpdateBotSessionRequest.
        // This request may take some time, and holding this lock prevents any
        // of the runner threads from starting work whilst we wait for the
        // request to return.
        sessionLock->unlock();

        grpc::CompletionQueue cq;

        // Add partial execution metadata (leaseID stored in the auxilary
        // metadata)
        for (const auto &[leaseID, metadata] : partialExecutionMetadata) {
            std::string header =
                "partial-execution-metadata-" + leaseID + "-bin";
            context.AddMetadata(header, metadata.SerializeAsString());
        }
        requestMetadata.attach_request_metadata(&context);

        std::unique_ptr<
            grpc::ClientAsyncResponseReaderInterface<proto::BotSession>>
            reader_ptr =
                stub->AsyncUpdateBotSession(&context, updateRequest, &cq);
        grpc::Status status;
        reader_ptr->Finish(&tempSession, &status, nullptr);
        waitForResponseOrCancel(&context, &cq, &status, cancelCall);

        // The request is over, we need to re-lock the session mutex before
        // handling the response
        sessionLock->lock();
        if (status.ok()) {
            // If the call was cancelled between receiving the response and
            // locking the mutex, return a cancelled status rather than
            // overwriting the potentially modified session with the response
            if (cancelCall) {
                return grpc::Status(
                    grpc::CANCELLED,
                    "Cancellation requested whilst receiving response");
            }

            if (executeOperationMetadataEntries != nullptr) {
                *executeOperationMetadataEntries =
                    extractExecuteOperationMetadata(
                        context.GetServerTrailingMetadata());
            }
            session->CopyFrom(tempSession);
            BUILDBOX_LOG_DEBUG("Bot session updated."
                               << sessionDebugInfo(session));
        }
        return status;
    };

    { // Timed Block
        buildboxcommon::buildboxcommonmetrics::MetricGuard<
            buildboxcommon::buildboxcommonmetrics::DurationMetricTimer>
            mt(MetricNames::TIMER_NAME_UPDATE_BOTSESSION);

        const int retryLimit = std::stoi(botsServerConnection.d_retryLimit);
        const int retryDelay = std::stoi(botsServerConnection.d_retryDelay);
        const std::chrono::seconds requestTimeout = std::chrono::seconds(
            std::stoi(botsServerConnection.d_requestTimeout));

        GrpcRetrier retrier(retryLimit, std::chrono::milliseconds(retryDelay),
                            updateInvocation, "UpdateBotSession()", {},
                            nullptr, requestTimeout);

        retrier.addOkStatusCode(grpc::StatusCode::CANCELLED);
        retrier.issueRequest();
        return retrier.status();
    }
}

grpc::Status BotSessionUtils::createBotSession(
    std::shared_ptr<proto::Bots::StubInterface> stub,
    const buildboxcommon::ConnectionOptions &botsServerConnection,
    proto::BotSession *session,
    const buildboxcommon::GrpcRetrier::GrpcStatusCodes &statusCodes,
    const buildboxcommon::RequestMetadataGenerator &requestMetadata)
{
    std::atomic_bool shouldCancel(false);
    return createBotSession(std::move(stub), botsServerConnection, session,
                            statusCodes, requestMetadata, shouldCancel);
}

grpc::Status BotSessionUtils::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)
{
    grpc::ClientContext context;
    proto::CreateBotSessionRequest createRequest;

    const auto updateInvocation = [&](grpc::ClientContext &context) {
        requestMetadata.attach_request_metadata(&context);
        BUILDBOX_LOG_DEBUG("Setting parent");
        createRequest.set_parent(botsServerConnection.d_instanceName);
        *createRequest.mutable_bot_session() = *session;
        BUILDBOX_LOG_DEBUG("Setting session");
        // If cancelCall is already set, return cancelled right away
        // instead of sending the request.
        if (cancelCall) {
            return grpc::Status(
                grpc::CANCELLED,
                "Cancellation requested before sending request");
        }

        grpc::CompletionQueue cq;
        std::unique_ptr<
            grpc::ClientAsyncResponseReaderInterface<proto::BotSession>>
            reader_ptr =
                stub->AsyncCreateBotSession(&context, createRequest, &cq);
        grpc::Status status;
        reader_ptr->Finish(session, &status, nullptr);
        waitForResponseOrCancel(&context, &cq, &status, cancelCall);
        return status;
    };

    grpc::Status status;
    { // Timed Block
        buildboxcommon::buildboxcommonmetrics::MetricGuard<
            buildboxcommon::buildboxcommonmetrics::DurationMetricTimer>
            mt(MetricNames::TIMER_NAME_CREATE_BOTSESSION);

        const int retryLimit = std::stoi(botsServerConnection.d_retryLimit);
        const int retryDelay = std::stoi(botsServerConnection.d_retryDelay);
        const std::chrono::seconds requestTimeout = std::chrono::seconds(
            std::stoi(botsServerConnection.d_requestTimeout));

        GrpcRetrier retrier(retryLimit, std::chrono::milliseconds(retryDelay),
                            updateInvocation, "CreateBotSession()", {},
                            nullptr, requestTimeout);
        retrier.addOkStatusCode(grpc::StatusCode::CANCELLED);
        retrier.issueRequest();
        return retrier.status();
    }
}

BotSessionUtils::ExecuteOperationMetadataEntriesMultiMap
BotSessionUtils::extractExecuteOperationMetadata(
    const GrpcServerMetadataMultiMap &metadata)
{
    const auto executeOperationMetadataName = "executeoperationmetadata-bin";

    ExecuteOperationMetadataEntriesMultiMap res;
    for (const auto &entry : metadata) {
        if (entry.first == executeOperationMetadataName) {
            ExecuteOperationMetadata message;
            if (message.ParseFromString(
                    std::string(entry.second.data(), entry.second.length()))) {
                res.emplace(message.action_digest(), message);
            }
        }
    }
    return res;
}

void BotSessionUtils::waitForResponseOrCancel(
    grpc::ClientContext *context, grpc::CompletionQueue *cq,
    grpc::Status *status, const std::atomic_bool &cancelCall)
{
    void *tag = nullptr;
    bool ok = false;
    while (true) {
        const auto deadline =
            std::chrono::system_clock::now() + std::chrono::milliseconds(250);
        const auto asyncStatus = cq->AsyncNext(&tag, &ok, deadline);
        if (asyncStatus != grpc::CompletionQueue::TIMEOUT) {
            // Either we got an event on the queue or the
            // completion queue was shutdown. Both mean
            // we break out of the loop
            break;
        }
        if (cancelCall) {
            context->TryCancel();
            buildboxcommon::GrpcClient::ShutdownAndEmptyQueue(cq);
        }
    }
}

} // namespace buildboxworker
