/*
 * 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, Execute)
{
    // Job should be executed locally and no calls to testRemoteServer will be
    // made, so a fallback to local execution will not be made since it is
    // executed locally
    const auto actionDigest = generateSimpleShellAction("echo hello, world");

    std::atomic_bool stopRequested = false;
    ActionResult actionResult =
        execClient->executeAction(actionDigest, stopRequested);
    EXPECT_EQ(actionResult.exit_code(), 0);
}

TEST_F(HybridFallbackExecutionServiceTestFixture, ExecuteMissingAction)
{
    // Job will be executed locally and no calls to testRemoteServer will be
    // made and it will fail with grpc::StatusCode::NOT_FOUND
    const auto actionDigest =
        buildboxcommon::DigestGenerator::hash("unavailable");

    std::atomic_bool stopRequested = false;
    try {
        execClient->executeAction(actionDigest, stopRequested);
        // If no exception is thrown, fail the test
        FAIL() << "Expected GrpcError exception to be thrown.";
    }
    catch (GrpcError &e) {
        EXPECT_TRUE(e.status.error_code() == grpc::StatusCode::NOT_FOUND);
    }
    ASSERT_THROW(execClient->executeAction(actionDigest, stopRequested),
                 GrpcError);
}

TEST_F(HybridFallbackExecutionServiceTestFixture, GetOperationAndWaitExecution)
{
    // Job will be executed locally and no calls to testRemoteServer will be
    // made since it is only one job
    const auto actionDigest = generateSimpleShellAction("sleep 1");

    std::atomic_bool stopRequested = false;
    auto operation =
        execClient->asyncExecuteAction(actionDigest, stopRequested);
    EXPECT_NE(operation.name(), "");
    EXPECT_FALSE(operation.done());

    operation = execClient->getOperation(operation.name());
    EXPECT_FALSE(operation.done());

    operation = execClient->waitExecution(operation.name());
    expectOperationSuccessful(operation);
}

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

    // Job will be executed locally and no calls to testRemoteServer will be
    // made since it is only one job
    const auto actionDigest = generateSimpleShellAction("sleep 1");

    std::atomic_bool stopRequested = false;
    auto operation =
        execClient->asyncExecuteAction(actionDigest, stopRequested);
    EXPECT_NE(operation.name(), "");
    EXPECT_FALSE(operation.done());

    while (!operation.done()) {
        operation = execClient->getOperation(operation.name());
        std::this_thread::sleep_for(
            std::chrono::milliseconds(POLLING_INTERVAL_MS));
    }

    expectOperationSuccessful(operation);
}

TEST_F(HybridFallbackExecutionServiceTestFixture, CancelOperation)
{
    // Job will be executed locally and no calls to testRemoteServer will be
    // made since it is only one job
    const auto actionDigest = generateSimpleShellAction("sleep 10000");

    std::atomic_bool stopRequested = false;
    auto operation =
        execClient->asyncExecuteAction(actionDigest, stopRequested);
    EXPECT_NE(operation.name(), "");
    EXPECT_FALSE(operation.done());

    execClient->cancelOperation(operation.name());

    operation = execClient->waitExecution(operation.name());

    expectOperationCancelled(operation);
}

TEST_F(HybridFallbackExecutionServiceTestFixture, Queue)
{
    const auto actionDigest1 =
        generateSimpleShellAction("sleep 10000 && echo 1");
    const auto actionDigest2 =
        generateSimpleShellAction("sleep 10000 && echo 2");
    const auto actionDigest3 = generateSimpleShellAction("echo hello, world");

    std::thread serverHandler([this, actionDigest3]() {
        auto operation = getExpectedOperation();
        operation.set_done(false);
        GrpcTestServerContext ctx(
            &testRemoteServer,
            "/build.bazel.remote.execution.v2.Execution/Execute");
        ctx.read(getExpectedExecuteRequest(actionDigest3));
        ctx.writeAndFinish(operation);
        completeOperation(&operation);
        // Handle expected WaitExecution request
        GrpcTestServerContext ctx2(
            &testRemoteServer,
            "/build.bazel.remote.execution.v2.Execution/WaitExecution");
        ctx2.read(getExpectedWaitExecutionRequest(operation));
        ctx2.writeAndFinish(operation);
    });
    try {

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

        // The third action should complete via remote execution despite the
        // first two (local) actions taking forever.
        operation3 = execClient->waitExecution(operation3.name());

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

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

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