/*
 * 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, QueueRemoteFail)
{
    constexpr int TEST_TIME_SLICE = 4;
    constexpr int TASK_DURATION_SECONDS = TEST_TIME_SLICE;     // 4s
    constexpr int MIN_EXECUTION_SECONDS = TEST_TIME_SLICE;     // 4s
    constexpr int MAX_EXECUTION_SECONDS = TEST_TIME_SLICE * 2; // 8s

    const auto start = std::chrono::steady_clock::now();

    const auto actionDigest1 = generateSimpleShellAction(
        "sleep " + std::to_string(TASK_DURATION_SECONDS) + " && echo 1");
    const auto actionDigest2 = generateSimpleShellAction(
        "sleep " + std::to_string(TASK_DURATION_SECONDS) + " && echo 2");
    const auto actionDigest3 = generateSimpleShellAction(
        "sleep " + std::to_string(TASK_DURATION_SECONDS) + " && echo 3");

    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);
        // Sleep to fake job execution
        sleep(TASK_DURATION_SECONDS);
        internalErrorOperationExecuteResponse(&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);

        operation1 = execClient->waitExecution(operation1.name());
        operation2 = execClient->waitExecution(operation2.name());
        operation3 = execClient->waitExecution(operation3.name());
        expectOperationSuccessful(operation1);
        expectOperationSuccessful(operation2);
        // WaitExecution will fail with INTERNAL and will not fallback
        expectOperationExitCode(operation3, google::rpc::INTERNAL);

        const auto end = std::chrono::steady_clock::now();

        // Verify execution took at least as long as individual tasks
        EXPECT_GT(end - start, std::chrono::seconds(MIN_EXECUTION_SECONDS));
        // Verify parallel execution kept total time reasonable
        EXPECT_LT(end - start, std::chrono::seconds(MAX_EXECUTION_SECONDS));

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

TEST_F(HybridFallbackExecutionServiceTestFixture, QueueRemoteFailWithException)
{
    constexpr int TEST_TIME_SLICE = 4;
    constexpr int TASK_DURATION_SECONDS = TEST_TIME_SLICE;       // 4s
    constexpr int FIRST_BATCH_MIN_SECONDS = TEST_TIME_SLICE;     // 4s
    constexpr int FIRST_BATCH_MAX_SECONDS = TEST_TIME_SLICE * 2; // 8s
    constexpr int FALLBACK_MIN_SECONDS = TEST_TIME_SLICE * 2;    // 8s
    constexpr int FALLBACK_MAX_SECONDS = TEST_TIME_SLICE * 3;    // 12s

    const auto start = std::chrono::steady_clock::now();

    const auto actionDigest1 = generateSimpleShellAction(
        "sleep " + std::to_string(TASK_DURATION_SECONDS) + " && echo 1");
    const auto actionDigest2 = generateSimpleShellAction(
        "sleep " + std::to_string(TASK_DURATION_SECONDS) + " && echo 2");
    const auto actionDigest3 = generateSimpleShellAction(
        "sleep " + std::to_string(TASK_DURATION_SECONDS) + " && echo 3");
    const auto actionDigest4 = generateSimpleShellAction(
        "sleep " + std::to_string(TASK_DURATION_SECONDS) + " && 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 abortedStatus(
            grpc::StatusCode::ABORTED,
            "Simulated ABORTED error during creation of the Operation");
        ctx2.finish(abortedStatus);
    });

    try {
        std::atomic_bool stopRequested = false;
        auto operation1 =
            execClient->asyncExecuteAction(actionDigest1, stopRequested);
        auto operation2 =
            execClient->asyncExecuteAction(actionDigest2, stopRequested);
        auto operation3 =
            execClient->asyncExecuteAction(actionDigest3, stopRequested);
        try {
            auto operation4 =
                execClient->asyncExecuteAction(actionDigest4, stopRequested);
        }
        catch (GrpcError &e) {
            EXPECT_TRUE(e.status.error_code() == grpc::StatusCode::ABORTED);
        }

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

        const auto first2end = std::chrono::steady_clock::now();
        // It should fallback and execute locally even if it queues up more
        // than the queue size
        operation3 = execClient->waitExecution(operation3.name());

        const auto nextEnd = std::chrono::steady_clock::now();

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

        // Verify first two operations completed in parallel within expected
        // timeframe
        EXPECT_GT(first2end - start,
                  std::chrono::seconds(FIRST_BATCH_MIN_SECONDS));
        EXPECT_LT(first2end - start,
                  std::chrono::seconds(FIRST_BATCH_MAX_SECONDS));
        // Verify fallback operation completed within extended timeframe
        EXPECT_GT(nextEnd - start, std::chrono::seconds(FALLBACK_MIN_SECONDS));
        EXPECT_LT(nextEnd - start, std::chrono::seconds(FALLBACK_MAX_SECONDS));

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