/*
 * Copyright 2024 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_LOCALEXECUTIONSCHEDULER_H
#define INCLUDED_BUILDBOXCASD_LOCALEXECUTIONSCHEDULER_H

#include <buildboxcommon_connectionoptions.h>
#include <buildboxcommon_localexecutionclient.h>

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <deque>
#include <memory>
#include <string>
#include <thread>
#include <vector>

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

using grpc::ServerContext;
using grpc::ServerWriterInterface;
using grpc::Status;

namespace buildboxcasd {

class LocalExecutionScheduler final {
  public:
    explicit LocalExecutionScheduler(
        const std::string &bindAddress, const std::string &runnerCommand,
        const std::vector<std::string> &extraRunArgs, int maxJobs);

    virtual ~LocalExecutionScheduler();

    void setCacheFailures(bool cacheFailures);

    std::shared_ptr<buildboxcommon::LocalExecutionClient>
    createClient(const std::string &casInstanceNameForRunner);

    grpc::Status
    Execute(ServerContext *ctx, const ExecuteRequest &request,
            ServerWriterInterface<Operation> *writer,
            std::shared_ptr<buildboxcommon::LocalExecutionClient> &execClient,
            const std::string &operationPrefix,
            std::optional<int> queueLimit = std::nullopt);

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

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

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

    void stop();

    LocalExecutionScheduler(const LocalExecutionScheduler &) = delete;
    LocalExecutionScheduler &
    operator=(const LocalExecutionScheduler &) = delete;
    LocalExecutionScheduler(LocalExecutionScheduler &&) = delete;
    LocalExecutionScheduler &operator=(LocalExecutionScheduler &&) = delete;

  private:
    struct InternalOperation {
        std::string name;
        std::shared_ptr<buildboxcommon::LocalExecutionClient> execClient;
        ExecuteRequest executeRequest;
        std::mutex mutex;
        Operation operation;
        std::condition_variable cv;
        bool queued = true;
        bool done = false;
    };

    void workerThread();

    grpc::Status
    waitExecutionInternal(ServerContext *ctx, const std::string &name,
                          const std::shared_ptr<InternalOperation> &op,
                          ServerWriterInterface<Operation> *writer);

    std::string d_bindAddress;
    std::string d_runnerCommand;
    std::vector<std::string> d_extraRunArgs;
    bool d_cacheFailures = false;

    std::vector<std::thread> d_workerThreads;

    std::mutex d_mutex;
    std::condition_variable d_cv;
    size_t d_workerThreadsWaiting = 0;
    std::atomic_bool d_stopRequested = false;
    std::deque<std::shared_ptr<InternalOperation>> d_queue;
    std::map<std::string, std::shared_ptr<InternalOperation>> d_operationMap;
    std::deque<std::pair<std::string, std::chrono::system_clock::time_point>>
        d_doneOperations;
};

} // namespace buildboxcasd

#endif // INCLUDED_BUILDBOXCASD_LOCALEXECUTIONSCHEDULER_H
