/*
 * Copyright 2022 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 <buildboxcasd_executionproxyinstance.h>
#include <buildboxcasd_requestcontextmanager.h>
#include <buildboxcommon_logging.h>

#include <utility>

using namespace buildboxcasd;
using namespace buildboxcommon;

ExecutionProxyInstance::ExecutionProxyInstance(
    const buildboxcommon::ConnectionOptions &execution_endpoint,
    const std::string &instance_name,
    const std::shared_ptr<LocalExecutionScheduler> &localScheduler,
    const std::string &casInstanceNameForRunner,
    std::optional<int> hybridQueueLimit,
    const bool enableFallbackToLocalExecution)
    : ExecutionInstance(instance_name),
      d_grpc_client(std::make_shared<GrpcClient>()),
      d_remote_instance_name(execution_endpoint.d_instanceName),
      d_stop_requested(false), d_localScheduler(localScheduler),
      d_casInstanceNameForRunner(casInstanceNameForRunner),
      d_hybridQueueLimit(hybridQueueLimit),
      d_enableFallbackToLocalExecution(enableFallbackToLocalExecution)
{
    d_grpc_client->init(execution_endpoint);
    RequestContextManager::configureGrpcClient(d_grpc_client.get());
    d_exec_client =
        std::make_shared<RemoteExecutionClient>(d_grpc_client, nullptr);
    d_exec_client->init();

    if (d_localScheduler &&
        (d_hybridQueueLimit.has_value() || d_enableFallbackToLocalExecution)) {
        d_localExecClient =
            d_localScheduler->createClient(casInstanceNameForRunner);
    }
    else if (d_hybridQueueLimit.has_value() ||
             d_enableFallbackToLocalExecution) {
        BUILDBOX_LOG_ERROR("Local execution scheduler is required for hybrid "
                           "or fallback execution");
        throw std::runtime_error("Local execution scheduler is required for "
                                 "hybrid or fallback execution");
    }
}

ExecutionProxyInstance::ExecutionProxyInstance(
    const ExecutionProxyInstance &other, const std::string &instanceName)
    : ExecutionInstance(instanceName), d_grpc_client(other.d_grpc_client),
      d_exec_client(other.d_exec_client),
      d_remote_instance_name(other.d_remote_instance_name),
      d_stop_requested(false), d_localScheduler(other.d_localScheduler),
      d_casInstanceNameForRunner(other.d_casInstanceNameForRunner),
      d_hybridQueueLimit(other.d_hybridQueueLimit),
      d_enableFallbackToLocalExecution(other.d_enableFallbackToLocalExecution)
{
    /* Don't share LocalExecutionClient with the other instance as `stop()`
     * in this instance must not cancel operations of the other instance.
     */
    if (other.d_localExecClient) {
        d_localExecClient =
            d_localScheduler->createClient(d_casInstanceNameForRunner);
    }
}

std::shared_ptr<ExecutionInstance>
ExecutionProxyInstance::clone(const std::string &instanceName)
{
    return std::make_shared<ExecutionProxyInstance>(*this, instanceName);
}

grpc::Status
ExecutionProxyInstance::Execute(ServerContext *ctx,
                                const ExecuteRequest &request,
                                ServerWriterInterface<Operation> *writer)
{
    if (request.action_digest().hash().empty()) {
        return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                            "Missing action digest");
    }

    std::string operationPrefix = d_instance_name + "#";

    if (d_localExecClient && d_hybridQueueLimit.has_value()) {
        auto status =
            d_localScheduler->Execute(ctx, request, writer, d_localExecClient,
                                      operationPrefix, d_hybridQueueLimit);
        if (status.error_code() != grpc::RESOURCE_EXHAUSTED) {
            return status;
        }
    }

    std::function<bool()> stop_lambda = [&] {
        return d_stop_requested.load() || ctx->IsCancelled();
    };
    try {
        ExecuteRequest remote_request;
        remote_request.CopyFrom(request);
        remote_request.set_instance_name(d_remote_instance_name);
        return d_exec_client->proxyExecuteRequest(remote_request, stop_lambda,
                                                  writer, operationPrefix);
    }
    catch (GrpcError &e) {
        if (d_localExecClient && d_enableFallbackToLocalExecution &&
            fallbackErrorCodes.find(e.status.error_code()) !=
                fallbackErrorCodes.end()) {
            BUILDBOX_LOG_WARNING("Failed to execute request remotely with: "
                                 << e.what()
                                 << ". Queuing request for local execution");
            // Execute locally ignoring the queue limit. Calling Execute
            // without a queue limit will result in the job executing without
            // the queue limit being enforced
            return d_localScheduler->Execute(
                ctx, request, writer, d_localExecClient, operationPrefix);
        }
        return e.status;
    }
}

grpc::Status
ExecutionProxyInstance::WaitExecution(ServerContext *ctx,
                                      const WaitExecutionRequest &request,
                                      ServerWriterInterface<Operation> *writer)
{
    if (d_localExecClient) {
        auto status = d_localScheduler->WaitExecution(ctx, request, writer);
        if (status.error_code() != grpc::NOT_FOUND) {
            return status;
        }
    }

    std::function<bool()> stop_lambda = [&] {
        return d_stop_requested.load() || ctx->IsCancelled();
    };
    try {
        WaitExecutionRequest remote_request;
        remote_request.CopyFrom(request);
        remote_request.set_name(rewriteOperationName(request.name()));
        std::string operation_prefix = d_instance_name + "#";
        return d_exec_client->proxyWaitExecutionRequest(
            remote_request, stop_lambda, writer, operation_prefix);
    }
    catch (GrpcError &e) {
        return e.status;
    }
}

grpc::Status
ExecutionProxyInstance::GetOperation(const GetOperationRequest &request,
                                     Operation *response)
{
    if (d_localExecClient) {
        auto status = d_localScheduler->GetOperation(request, response);
        if (status.error_code() != grpc::NOT_FOUND) {
            return status;
        }
    }

    try {
        GetOperationRequest remote_request;
        remote_request.CopyFrom(request);
        remote_request.set_name(rewriteOperationName(request.name()));
        grpc::Status status =
            d_exec_client->proxyGetOperationRequest(remote_request, response);
        response->set_name(d_instance_name + "#" + response->name());
        return status;
    }
    catch (GrpcError &e) {
        return e.status;
    }
}

grpc::Status
ExecutionProxyInstance::ListOperations(const ListOperationsRequest &request,
                                       ListOperationsResponse *response)
{
    if (d_localExecClient) {
        return grpc::Status(grpc::StatusCode::UNIMPLEMENTED, __func__);
    }

    try {
        ListOperationsRequest remote_request;
        remote_request.CopyFrom(request);
        remote_request.set_name(d_remote_instance_name);
        grpc::Status status = d_exec_client->proxyListOperationsRequest(
            remote_request, response);

        // Rewrite operation names to include the local instance name.
        for (Operation &operation : *response->mutable_operations()) {
            operation.set_name(d_instance_name + "#" + operation.name());
        }

        return status;
    }
    catch (GrpcError &e) {
        return e.status;
    }
}

grpc::Status
ExecutionProxyInstance::CancelOperation(const CancelOperationRequest &request,
                                        google::protobuf::Empty *response)
{
    if (d_localExecClient) {
        auto status = d_localScheduler->CancelOperation(request, response);
        if (status.error_code() != grpc::NOT_FOUND) {
            return status;
        }
    }

    try {
        CancelOperationRequest remote_request;
        remote_request.CopyFrom(request);
        remote_request.set_name(rewriteOperationName(request.name()));
        return d_exec_client->proxyCancelOperationRequest(remote_request,
                                                          response);
    }
    catch (GrpcError &e) {
        return e.status;
    }
}

void ExecutionProxyInstance::stop()
{
    d_stop_requested = true;

    if (d_localExecClient) {
        d_localExecClient->cancelAllOperations();
    }
}

std::string ExecutionProxyInstance::rewriteOperationName(std::string name)
{
    size_t idx = name.find_first_of('#');
    if (idx == std::string::npos) {
        return name;
    }
    return name.substr(idx + 1);
}
