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

#include <build/bazel/remote/execution/v2/remote_execution.pb.h>
#include <buildboxcasd_fslocalactionstorage.h>
#include <buildboxcasd_fslocalassetstorage.h>
#include <buildboxcommon_grpctestserver.h>
#include <google/longrunning/operations.pb.h>

#include <buildboxcasd_localexecutioninstance.h>
#include <buildboxcasd_server.h>
#include <buildboxcommon_fslocalcas.h>

#include <buildboxcommon_casclient.h>
#include <buildboxcommon_digestgenerator.h>
#include <buildboxcommon_grpcclient.h>
#include <buildboxcommon_protos.h>
#include <buildboxcommon_remoteexecutionclient.h>

#include <google/rpc/code.pb.h>
#include <gtest/gtest.h>
#include <unistd.h>

using namespace buildboxcasd;
using namespace buildboxcommon;
using namespace google::longrunning;

const auto digestFunctionInitializer = []() {
    buildboxcommon::DigestGenerator::init();
    return 0;
}();

TEST_F(HybridFallbackExecutionServiceTestFixture,
       TestFallbackResourceExhausted)
{
    constexpr int POLLING_INTERVAL_MS = 100;

    // Test servers are configured for 2 parallel jobs local and 1 remote job

    TemporaryDirectory tempDir;

    const auto actionDigest1 = generateSimpleShellAction(
        "touch " + tempDir.strname() + "/1 && sleep 10000");
    const auto actionDigest2 = generateSimpleShellAction(
        "touch " + tempDir.strname() + "/2 && sleep 10000");
    const auto actionDigest3 = generateSimpleShellAction(
        "touch " + tempDir.strname() + "/3 && sleep 10000");
    const auto actionDigest4 = generateSimpleShellAction("echo hello, world");

    std::thread serverHandler([this, actionDigest3, actionDigest4]() {
        GrpcTestServerContext ctx(
            &testRemoteServer,
            "/build.bazel.remote.execution.v2.Execution/Execute");
        ctx.read(getExpectedExecuteRequest(actionDigest3));
        grpc::Status errorStatus(
            grpc::StatusCode::RESOURCE_EXHAUSTED,
            "Simulated error during creation of the Operation");
        ctx.finish(errorStatus);

        GrpcTestServerContext ctx2(
            &testRemoteServer,
            "/build.bazel.remote.execution.v2.Execution/Execute");
        ctx2.read(getExpectedExecuteRequest(actionDigest4));
        ctx2.finish(errorStatus);
    });

    try {
        std::atomic_bool stopRequested = false;
        auto operation1 =
            execClient->asyncExecuteAction(actionDigest1, stopRequested);
        auto operation2 =
            execClient->asyncExecuteAction(actionDigest2, stopRequested);
        auto operation3 =
            execClient->asyncExecuteAction(actionDigest3, stopRequested);
        auto operation4 =
            execClient->asyncExecuteAction(actionDigest4, stopRequested);

        // Make sure actions have actually started before cancelling them
        // (ensures that we get timestamps needed for verification)
        while (!FileUtils::existsInDir(tempDir.fd(), "1") ||
               !FileUtils::existsInDir(tempDir.fd(), "2")) {
            std::this_thread::sleep_for(
                std::chrono::milliseconds(POLLING_INTERVAL_MS));
        }

        execClient->cancelOperation(operation1.name());
        operation1 = execClient->waitExecution(operation1.name());
        execClient->cancelOperation(operation2.name());
        operation2 = execClient->waitExecution(operation2.name());

        // It should fallback and execute locally even if it queues up more
        // than the queue size
        operation4 = execClient->waitExecution(operation4.name());

        while (!FileUtils::existsInDir(tempDir.fd(), "3")) {
            std::this_thread::sleep_for(
                std::chrono::milliseconds(POLLING_INTERVAL_MS));
        }
        execClient->cancelOperation(operation3.name());
        operation3 = execClient->waitExecution(operation3.name());

        const auto execMetadata1 = expectOperationCancelled(operation1);
        const auto execMetadata2 = expectOperationCancelled(operation2);
        const auto execMetadata3 = expectOperationCancelled(operation3);
        const auto execMetadata4 = expectOperationSuccessful(operation4);

        // Verify that the first two operations were running in parallel
        EXPECT_LT(execMetadata1.worker_start_timestamp(),
                  execMetadata2.worker_completed_timestamp());
        EXPECT_LT(execMetadata2.worker_start_timestamp(),
                  execMetadata1.worker_completed_timestamp());

        // Verify that operation 4 was started after operation 1 was cancelled
        EXPECT_GT(execMetadata4.worker_start_timestamp(),
                  execMetadata1.worker_completed_timestamp());

        EXPECT_FALSE(isRemoteOperation(operation1));
        EXPECT_FALSE(isRemoteOperation(operation2));
        EXPECT_FALSE(isRemoteOperation(operation3));
        EXPECT_FALSE(isRemoteOperation(operation4));
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(HybridFallbackExecutionServiceTestFixture,
       HybridFallbackCancellationTest)
{
    TemporaryDirectory tempDir;

    const auto actionDigest1 =
        generateSimpleShellAction("while ! test -e " + tempDir.strname() +
                                  "/run; do sleep 1; done; echo 1");
    const auto actionDigest2 =
        generateSimpleShellAction("while ! test -e " + tempDir.strname() +
                                  "/run; do sleep 1; done; echo 2");
    const auto actionDigest3 =
        generateSimpleShellAction("sleep 10000 && echo 3");
    const auto actionDigest4 =
        generateSimpleShellAction("sleep 10000 && echo 4");

    std::thread serverHandler([this, actionDigest3, actionDigest4]() {
        auto operation = getExpectedOperation();
        operation.set_done(false);
        GrpcTestServerContext ctx(
            &testRemoteServer,
            "/build.bazel.remote.execution.v2.Execution/Execute");
        ctx.read(getExpectedExecuteRequest(actionDigest3));
        grpc::Status errorStatus(grpc::StatusCode::RESOURCE_EXHAUSTED,
                                 "Simulated RESOURCE_EXHAUSTED error during "
                                 "creation of the Operation");
        ctx.finish(errorStatus);

        GrpcTestServerContext ctx2(
            &testRemoteServer,
            "/build.bazel.remote.execution.v2.Execution/Execute");
        ctx2.read(getExpectedExecuteRequest(actionDigest4));
        grpc::Status errorStatus2(grpc::StatusCode::RESOURCE_EXHAUSTED,
                                  "Simulated RESOURCE_EXHAUSTED error during "
                                  "creation of the Operation");
        ctx2.finish(errorStatus2);
    });

    try {
        std::atomic_bool stopRequested = false;
        auto operation1 =
            execClient->asyncExecuteAction(actionDigest1, stopRequested);
        auto operation2 =
            execClient->asyncExecuteAction(actionDigest2, stopRequested);
        auto operation3 =
            execClient->asyncExecuteAction(actionDigest3, stopRequested);
        auto operation4 =
            execClient->asyncExecuteAction(actionDigest4, stopRequested);

        // Signal the first two actions that they are allowed to run now that
        // everything is queued
        FileUtils::writeFileAtomically(tempDir.strname() + "/run", "");

        operation1 = execClient->waitExecution(operation1.name());
        operation2 = execClient->waitExecution(operation2.name());

        execClient->cancelOperation(operation3.name());
        execClient->cancelOperation(operation4.name());

        operation3 = execClient->waitExecution(operation3.name());
        operation4 = execClient->waitExecution(operation4.name());

        expectOperationSuccessful(operation1);
        expectOperationSuccessful(operation2);
        expectOperationCancelled(operation3);
        expectOperationCancelled(operation4);

        EXPECT_FALSE(isRemoteOperation(operation1));
        EXPECT_FALSE(isRemoteOperation(operation2));
        EXPECT_FALSE(isRemoteOperation(operation3));
        EXPECT_FALSE(isRemoteOperation(operation4));
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}
