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

#ifndef INCLUDED_BUILDBOXCASD_LOCALEXECUTIONPROXYINSTANCE_H
#define INCLUDED_BUILDBOXCASD_LOCALEXECUTIONPROXYINSTANCE_H

#include <buildboxcasd_executioninstance.h>
#include <buildboxcasd_localexecutionscheduler.h>
#include <buildboxcommon_connectionoptions.h>
#include <buildboxcommon_grpcclient.h>
#include <buildboxcommon_remoteexecutionclient.h>

using namespace build::bazel::remote::execution::v2;
using namespace google::longrunning;

using grpc::Status;

namespace buildboxcasd {

class ExecutionProxyInstance final : public ExecutionInstance {
  public:
    explicit ExecutionProxyInstance(
        const buildboxcommon::ConnectionOptions &execution_endpoint,
        const std::string &instance_name,
        const std::shared_ptr<LocalExecutionScheduler> &localScheduler =
            nullptr,
        const std::string &casInstanceNameForRunner = "",
        std::optional<int> hybridQueueLimit = std::nullopt,
        const bool enableFallbackToLocalExecution = false);

    ExecutionProxyInstance(const ExecutionProxyInstance &other,
                           const std::string &instanceName);

    std::shared_ptr<ExecutionInstance>
    clone(const std::string &instanceName) override;

    grpc::Status Execute(ServerContext *ctx, const ExecuteRequest &request,
                         ServerWriterInterface<Operation> *writer) override;

    grpc::Status
    WaitExecution(ServerContext *ctx, const WaitExecutionRequest &request,
                  ServerWriterInterface<Operation> *writer) override;

    grpc::Status GetOperation(const GetOperationRequest &request,
                              Operation *response) override;

    grpc::Status ListOperations(const ListOperationsRequest &request,
                                ListOperationsResponse *response) override;

    grpc::Status CancelOperation(const CancelOperationRequest &request,
                                 google::protobuf::Empty *response) override;

    void stop() override;

    std::string rewriteOperationName(std::string name);

  private:
    std::shared_ptr<buildboxcommon::GrpcClient> d_grpc_client;
    std::shared_ptr<buildboxcommon::RemoteExecutionClient> d_exec_client;
    std::string d_remote_instance_name;
    std::atomic_bool d_stop_requested;

    std::shared_ptr<LocalExecutionScheduler> d_localScheduler;
    std::shared_ptr<buildboxcommon::LocalExecutionClient> d_localExecClient;
    std::string d_casInstanceNameForRunner;
    std::optional<int> d_hybridQueueLimit;
    bool d_enableFallbackToLocalExecution;
    const std::unordered_set<grpc::StatusCode> fallbackErrorCodes = {
        grpc::StatusCode::RESOURCE_EXHAUSTED};
};

} // namespace buildboxcasd

#endif // INCLUDED_BUILDBOXCASD_LOCALEXECUTIONPROXYINSTANCE_H
