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

#include <buildboxcasd_localexecutionscheduler.h>
#include <buildboxcasd_requestcontextmanager.h>
#include <buildboxcommon_connectionoptions.h>

#include <buildboxcommon_logging.h>
#include <buildboxcommon_stringutils.h>

#include <google/rpc/code.pb.h>
#include <google/rpc/status.pb.h>

#include <memory>
#include <thread>

#define POLL_WAIT std::chrono::milliseconds(500)
#define OPERATION_TTL std::chrono::minutes(5)

using namespace buildboxcasd;
using namespace buildboxcommon;

LocalExecutionScheduler::LocalExecutionScheduler(
    const std::string &bindAddress, const std::string &runnerCommand,
    const std::vector<std::string> &extraRunArgs, int maxJobs)
    : d_bindAddress(bindAddress), d_runnerCommand(runnerCommand),
      d_extraRunArgs(extraRunArgs)
{
    for (int i = 0; i < maxJobs; i++) {
        d_workerThreads.emplace_back(&LocalExecutionScheduler::workerThread,
                                     this);
    }
}

LocalExecutionScheduler::~LocalExecutionScheduler()
{
    this->stop();

    for (auto &thread : d_workerThreads) {
        thread.join();
    }
}

void LocalExecutionScheduler::setCacheFailures(bool cacheFailures)
{
    this->d_cacheFailures = cacheFailures;
}

std::shared_ptr<LocalExecutionClient> LocalExecutionScheduler::createClient(
    const std::string &casInstanceNameForRunner)
{
    /*
     * ConnectionOptions is required for the buildbox-casd command line.
     * I.e., this can't be replaced with process-internal access to CAS
     * storage.
     */
    buildboxcommon::ConnectionOptions casdEndpoint;
    if (d_bindAddress.find("unix:") == 0) {
        casdEndpoint.setUrl(d_bindAddress);
    }
    else {
        casdEndpoint.setUrl("http://" + d_bindAddress);
    }
    casdEndpoint.setInstanceName(casInstanceNameForRunner);

    buildboxcommon::ConnectionOptions raEndpoint;
    if (d_bindAddress.find("unix:") == 0) {
        raEndpoint.setUrl(d_bindAddress);
    }
    else {
        raEndpoint.setUrl("http://" + d_bindAddress);
    }
    raEndpoint.setInstanceName(casInstanceNameForRunner);

    auto grpcClient = std::make_shared<buildboxcommon::GrpcClient>();
    grpcClient->init(casdEndpoint);
    RequestContextManager::configureGrpcClient(grpcClient.get());

    auto execClient = std::make_shared<LocalExecutionClient>(
        casdEndpoint, raEndpoint, grpcClient);
    execClient->setCacheFailures(d_cacheFailures);
    execClient->setRunner(d_runnerCommand, d_extraRunArgs);
    execClient->init();
    return execClient;
}

void LocalExecutionScheduler::workerThread()
{
    std::unique_lock lock(d_mutex, std::defer_lock);

    while (true) {
        lock.lock();
        d_workerThreadsWaiting++;
        d_cv.wait(lock, [&] { return d_stopRequested || !d_queue.empty(); });
        d_workerThreadsWaiting--;

        if (d_stopRequested) {
            break;
        }

        auto op = d_queue.front();
        d_queue.pop_front();
        lock.unlock();

        std::unique_lock opLock(op->mutex);
        op->queued = false;
        if (op->done) {
            // Operation was cancelled before starting execution
            continue;
        }

        try {
            op->operation = op->execClient->asyncExecuteAction(
                op->executeRequest.action_digest(), d_stopRequested,
                op->executeRequest.skip_cache_lookup());

            if (!op->operation.done()) {
                const auto operationName = op->operation.name();
                opLock.unlock();
                auto operation = op->execClient->waitExecution(operationName);
                opLock.lock();
                op->operation = operation;
            }
        }
        catch (GrpcError &e) {
            if (!opLock.owns_lock()) {
                // Reacquire lock if exception was thrown while unlocked
                opLock.lock();
            }

            ExecuteResponse executeResponse;
            google::rpc::Status *status = executeResponse.mutable_status();
            status->set_code(e.status.error_code());
            status->set_message(e.status.error_message());

            op->operation.set_done(true);
            op->operation.mutable_response()->PackFrom(executeResponse);
        }
        catch (const std::exception &e) {
            if (!opLock.owns_lock()) {
                // Reacquire lock if exception was thrown while unlocked
                opLock.lock();
            }

            ExecuteResponse executeResponse;
            google::rpc::Status *status = executeResponse.mutable_status();
            status->set_code(google::rpc::Code::INTERNAL);
            status->set_message(e.what());

            op->operation.set_done(true);
            op->operation.mutable_response()->PackFrom(executeResponse);
        }

        op->done = true;
        opLock.unlock();

        // Cleanup operations that have been completed for a while and then
        // register the just completed operation with a time-to-live.
        lock.lock();
        const auto now = std::chrono::system_clock::now();
        while (!d_doneOperations.empty()) {
            const auto &pair = d_doneOperations.front();
            if (pair.second < now) {
                d_operationMap.erase(pair.first);
                d_doneOperations.pop_front();
            }
            else {
                break;
            }
        }
        d_doneOperations.emplace_back(op->name, now + OPERATION_TTL);
        lock.unlock();

        // Wake everyone waiting for this operation to be done
        op->cv.notify_all();
    }
}

grpc::Status LocalExecutionScheduler::Execute(
    ServerContext *ctx, const ExecuteRequest &request,
    ServerWriterInterface<Operation> *writer,
    std::shared_ptr<buildboxcommon::LocalExecutionClient> &execClient,
    const std::string &operationPrefix, std::optional<int> queueLimit)
{
    Operation operation;
    operation.set_name(operationPrefix + StringUtils::getUUIDString());

    if (!request.skip_cache_lookup()) {
        ExecuteResponse executeResponse;
        if (execClient->fetchFromActionCache(
                request.action_digest(), {},
                executeResponse.mutable_result())) {
            executeResponse.set_cached_result(true);

            BUILDBOX_LOG_DEBUG("Action Cache hit for ["
                               << request.action_digest() << "]");

            operation.set_done(true);
            operation.mutable_response()->PackFrom(executeResponse);

            if (!writer->Write(operation)) {
                // The stream has been closed
                return grpc::Status::CANCELLED;
            }

            return grpc::Status::OK;
        }
    }

    auto op = std::make_shared<InternalOperation>();
    op->name = operation.name();
    op->execClient = execClient;
    op->executeRequest = request;
    {
        std::lock_guard lock(d_mutex);

        if (queueLimit.has_value() &&
            d_queue.size() + 1 > d_workerThreadsWaiting + queueLimit.value()) {
            return grpc::Status(grpc::StatusCode::RESOURCE_EXHAUSTED,
                                "Local execution queue above limit");
        }

        d_operationMap.emplace(operation.name(), op);
        d_queue.push_back(op);

        if (d_workerThreadsWaiting >= d_queue.size()) {
            // The action will start execution immediately.
            // Skip action cache lookup in worker thread.
            op->executeRequest.set_skip_cache_lookup(true);
        }
    }
    // Wake a worker thread to start execution
    d_cv.notify_one();

    if (!writer->Write(operation)) {
        // The stream has been closed
        return grpc::Status::CANCELLED;
    }

    return waitExecutionInternal(ctx, operation.name(), op, writer);
}

grpc::Status LocalExecutionScheduler::waitExecutionInternal(
    ServerContext *ctx, const std::string &name,
    const std::shared_ptr<InternalOperation> &op,
    ServerWriterInterface<Operation> *writer)
{
    std::unique_lock opLock(op->mutex);
    while (!op->done) {
        // Set a timeout to allow releasing this thread before execution
        // is complete if the client cancels the gRPC request.
        constexpr int TIMEOUT_SECONDS = 5;
        op->cv.wait_for(opLock, std::chrono::seconds(TIMEOUT_SECONDS));

        if (ctx->IsCancelled()) {
            return grpc::Status::CANCELLED;
        }
    }

    Operation operation = op->operation;
    operation.set_name(name);
    opLock.unlock();

    if (!writer->Write(operation)) {
        // The stream has been closed
        return grpc::Status::CANCELLED;
    }

    return grpc::Status::OK;
}

grpc::Status LocalExecutionScheduler::WaitExecution(
    ServerContext *ctx, const WaitExecutionRequest &request,
    ServerWriterInterface<Operation> *writer)
{
    std::shared_ptr<InternalOperation> op;
    {
        std::lock_guard lock(d_mutex);
        auto it = d_operationMap.find(request.name());
        if (it == d_operationMap.end()) {
            return grpc::Status(grpc::StatusCode::NOT_FOUND,
                                "The operation was not found");
        }
        op = it->second;
    }

    return waitExecutionInternal(ctx, request.name(), op, writer);
}

grpc::Status
LocalExecutionScheduler::GetOperation(const GetOperationRequest &request,
                                      Operation *response)
{
    std::shared_ptr<InternalOperation> op;
    {
        std::lock_guard lock(d_mutex);
        auto it = d_operationMap.find(request.name());
        if (it == d_operationMap.end()) {
            return grpc::Status(grpc::StatusCode::NOT_FOUND,
                                "The operation was not found");
        }
        op = it->second;
    }

    std::unique_lock opLock(op->mutex);
    *response = op->operation;
    response->set_name(request.name());

    return grpc::Status::OK;
}

grpc::Status
LocalExecutionScheduler::CancelOperation(const CancelOperationRequest &request,
                                         google::protobuf::Empty *response)
{
    std::shared_ptr<InternalOperation> op;
    {
        std::lock_guard lock(d_mutex);
        auto it = d_operationMap.find(request.name());
        if (it == d_operationMap.end()) {
            return grpc::Status(grpc::StatusCode::NOT_FOUND,
                                "The operation was not found");
        }
        op = it->second;
    }

    std::unique_lock opLock(op->mutex);
    if (op->queued) {
        // Execution hasn't started yet

        ExecuteResponse executeResponse;
        google::rpc::Status *status = executeResponse.mutable_status();
        status->set_code(google::rpc::Code::CANCELLED);
        status->set_message("The operation was cancelled");

        op->operation.set_done(true);
        op->operation.mutable_response()->PackFrom(executeResponse);

        op->done = true;
        opLock.unlock();

        // Wake everyone waiting for this operation to be done
        op->cv.notify_all();
    }
    else if (!op->done) {
        // Execution has already started, use LocalExecutionClient's
        // operation name to cancel.
        if (!op->execClient->cancelOperation(op->operation.name())) {
            return grpc::Status(grpc::StatusCode::INTERNAL,
                                "Failed to cancel operation");
        }
    }

    *response = google::protobuf::Empty();
    return grpc::Status::OK;
}

void LocalExecutionScheduler::stop()
{
    {
        std::lock_guard lock(d_mutex);
        d_stopRequested = true;
    }
    d_cv.notify_all();
}
