/*
 * 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_fslocalactionstorage.h>
#include <buildboxcasd_fslocalassetstorage.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 <chrono>
#include <google/rpc/code.pb.h>
#include <gtest/gtest.h>

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

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

/***
 * Fixture to support testing buildbox-casd as a local Execution service
 */
class LocalExecutionServiceTestFixture : public ::testing::Test {
  protected:
    buildboxcommon::TemporaryDirectory socket_directory;
    const std::string TEST_SERVER_ADDRESS;

    buildboxcommon::TemporaryDirectory storage_root_directory;
    std::shared_ptr<FsLocalCas> cas_storage;
    std::shared_ptr<FsLocalAssetStorage> asset_storage;
    std::shared_ptr<FsLocalActionStorage> action_storage;
    std::shared_ptr<Server> server;

    std::shared_ptr<grpc::Channel> channel;
    std::shared_ptr<Capabilities::StubInterface> capabilitiesStub;
    std::shared_ptr<ByteStream::StubInterface> byteStreamStub;
    std::shared_ptr<ContentAddressableStorage::StubInterface> casStub;
    std::shared_ptr<LocalContentAddressableStorage::StubInterface>
        localcasStub;
    std::shared_ptr<Execution::StubInterface> executionStub;
    std::shared_ptr<google::longrunning::Operations::StubInterface>
        operationsStub;
    std::shared_ptr<GrpcClient> grpcClient;
    std::shared_ptr<CASClient> casClient;
    std::shared_ptr<RemoteExecutionClient> execClient;

    const std::string instance_name = "test";

    LocalExecutionServiceTestFixture(bool cacheFailures = false)
        : TEST_SERVER_ADDRESS(
              "unix://" + std::string(socket_directory.name()) + "/casd.sock"),
          grpcClient(std::make_shared<buildboxcommon::GrpcClient>()),
          casClient(std::make_shared<CASClient>(grpcClient)),
          execClient(
              std::make_shared<RemoteExecutionClient>(grpcClient, nullptr))
    {
        cas_storage =
            std::make_shared<FsLocalCas>(storage_root_directory.name());
        asset_storage = std::make_shared<FsLocalAssetStorage>(
            storage_root_directory.name());
        action_storage = std::make_shared<FsLocalActionStorage>(
            storage_root_directory.name());

        const char *runnerCommand = getenv("BUILDBOX_RUN");
        EXPECT_NE(runnerCommand, nullptr);
        const std::vector<std::string> extraRunArgs = {};

        auto scheduler = std::make_shared<LocalExecutionScheduler>(
            TEST_SERVER_ADDRESS, runnerCommand, extraRunArgs, /* maxJobs */ 2);

        scheduler->setCacheFailures(cacheFailures);

        server = std::make_shared<Server>(cas_storage, asset_storage,
                                          action_storage, scheduler);
        server->addLocalServerInstance(instance_name);
        server->addListeningPort(TEST_SERVER_ADDRESS);
        server->start();

        channel = grpc::CreateChannel(TEST_SERVER_ADDRESS,
                                      grpc::InsecureChannelCredentials());
        byteStreamStub = ByteStream::NewStub(channel);
        casStub = ContentAddressableStorage::NewStub(channel);
        localcasStub = LocalContentAddressableStorage::NewStub(channel);
        executionStub = Execution::NewStub(channel);
        operationsStub = Operations::NewStub(channel);
        capabilitiesStub = Capabilities::NewStub(channel);
        grpcClient->setInstanceName(instance_name);
        casClient->init(byteStreamStub, casStub, localcasStub, nullptr);
        execClient->init(executionStub, nullptr, operationsStub);
    }

    Digest generateSimpleShellAction(const std::string &commandString)
    {
        Action action;
        // Empty input root
        action.mutable_input_root_digest()->CopyFrom(
            DigestGenerator::hash(""));
        Command command;
        command.add_arguments("/bin/sh");
        command.add_arguments("-c");
        command.add_arguments(commandString);

        const auto commandDigest = casClient->uploadMessage(command);
        action.mutable_command_digest()->CopyFrom(commandDigest);
        auto actionDigest = casClient->uploadMessage(action);
        return actionDigest;
    }

    void expectOperationCancelled(const Operation &operation)
    {
        EXPECT_TRUE(operation.done());

        ExecuteResponse executeResponse;
        ASSERT_TRUE(operation.response().UnpackTo(&executeResponse));

        EXPECT_EQ(executeResponse.status().code(),
                  google::rpc::Code::CANCELLED);
    }

    ExecutedActionMetadata
    expectOperationSuccessful(const Operation &operation)
    {
        EXPECT_TRUE(operation.done());

        ExecuteResponse executeResponse;
        EXPECT_TRUE(operation.response().UnpackTo(&executeResponse));

        EXPECT_EQ(executeResponse.status().code(), google::rpc::Code::OK);

        const ActionResult actionResult = executeResponse.result();
        EXPECT_EQ(actionResult.exit_code(), 0);

        return actionResult.execution_metadata();
    }

    void expectOperationFailed(const Operation &operation)
    {
        EXPECT_TRUE(operation.done());

        ExecuteResponse executeResponse;
        ASSERT_TRUE(operation.response().UnpackTo(&executeResponse));

        EXPECT_EQ(executeResponse.status().code(), google::rpc::Code::OK);

        const ActionResult actionResult = executeResponse.result();
        EXPECT_NE(actionResult.exit_code(), 0);
    }
};

class FailureCachingLocalExecutionServiceTestFixture
    : public LocalExecutionServiceTestFixture {
  protected:
    FailureCachingLocalExecutionServiceTestFixture()
        : LocalExecutionServiceTestFixture(true)
    {
    }
};

TEST_F(LocalExecutionServiceTestFixture, Execute)
{
    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(LocalExecutionServiceTestFixture, ExecuteMissingAction)
{
    const auto actionDigest = DigestGenerator::hash("unavailable");

    std::atomic_bool stopRequested = false;
    ASSERT_THROW(execClient->executeAction(actionDigest, stopRequested),
                 GrpcError);
}

TEST_F(LocalExecutionServiceTestFixture, GetOperationAndWaitExecution)
{
    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(LocalExecutionServiceTestFixture, GetOperation)
{
    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(100));
    }

    expectOperationSuccessful(operation);
}

TEST_F(LocalExecutionServiceTestFixture, CancelOperation)
{
    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(LocalExecutionServiceTestFixture, Queue)
{
    // Test server is configured for 2 parallel jobs

    const auto actionDigest1 = generateSimpleShellAction("sleep 10000");
    const auto actionDigest2 = generateSimpleShellAction("echo hello 2");
    const auto actionDigest3 = generateSimpleShellAction("echo hello 3");

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

    // With the parallelism, the second and third actions should complete
    // despite the first action not completing.
    operation2 = execClient->waitExecution(operation2.name());
    operation3 = execClient->waitExecution(operation3.name());

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

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

    // The second and third actions should not have been running in parallel as
    // parallelism is limited to 2 jobs.
    EXPECT_GT(execMetadata3.execution_start_timestamp(),
              execMetadata2.execution_completed_timestamp());
}

TEST_F(LocalExecutionServiceTestFixture, ExecuteCache)
{
    const auto actionDigest = generateSimpleShellAction("echo hello, world");

    std::atomic_bool stopRequested = false;
    auto operation =
        execClient->asyncExecuteAction(actionDigest, stopRequested);

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

    ExecuteResponse executeResponse;
    ASSERT_TRUE(operation.response().UnpackTo(&executeResponse));
    EXPECT_FALSE(executeResponse.cached_result());

    auto operation2 =
        execClient->asyncExecuteAction(actionDigest, stopRequested);
    expectOperationSuccessful(operation2);

    ExecuteResponse executeResponse2;
    ASSERT_TRUE(operation2.response().UnpackTo(&executeResponse2));
    EXPECT_TRUE(executeResponse2.cached_result());
}

TEST_F(LocalExecutionServiceTestFixture, ExecuteCacheFailureDisabled)
{
    const auto actionDigest = generateSimpleShellAction("false");

    std::atomic_bool stopRequested = false;
    auto operation =
        execClient->asyncExecuteAction(actionDigest, stopRequested);

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

    ExecuteResponse executeResponse;
    ASSERT_TRUE(operation.response().UnpackTo(&executeResponse));
    EXPECT_FALSE(executeResponse.cached_result());

    auto operation2 =
        execClient->asyncExecuteAction(actionDigest, stopRequested);
    EXPECT_FALSE(operation2.done());
    operation2 = execClient->waitExecution(operation2.name());
    expectOperationFailed(operation2);

    ExecuteResponse executeResponse2;
    ASSERT_TRUE(operation2.response().UnpackTo(&executeResponse2));
    EXPECT_FALSE(executeResponse2.cached_result());
}

TEST_F(FailureCachingLocalExecutionServiceTestFixture,
       ExecuteCacheFailureEnabled)
{
    const auto actionDigest = generateSimpleShellAction("false");

    std::atomic_bool stopRequested = false;
    auto operation =
        execClient->asyncExecuteAction(actionDigest, stopRequested);

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

    ExecuteResponse executeResponse;
    ASSERT_TRUE(operation.response().UnpackTo(&executeResponse));
    EXPECT_FALSE(executeResponse.cached_result());

    auto operation2 =
        execClient->asyncExecuteAction(actionDigest, stopRequested);
    expectOperationFailed(operation2);

    ExecuteResponse executeResponse2;
    ASSERT_TRUE(operation2.response().UnpackTo(&executeResponse2));
    EXPECT_TRUE(executeResponse2.cached_result());
}

TEST_F(LocalExecutionServiceTestFixture, TestCapabilities)
{
    GetCapabilitiesRequest request;
    request.set_instance_name(instance_name);

    ServerCapabilities response;

    grpc::ClientContext client_context;
    const grpc::Status response_status =
        capabilitiesStub->GetCapabilities(&client_context, request, &response);
    ASSERT_EQ(response_status.error_code(), grpc::StatusCode::OK);

    ASSERT_EQ(response.cache_capabilities().max_batch_total_size_bytes(), 0);

    ASSERT_EQ(response.cache_capabilities().digest_functions_size(), 1);
#ifndef __SANITIZE_ADDRESS__
    ASSERT_EQ(response.cache_capabilities().digest_functions(0),
              DigestGenerator::digestFunction());
#endif
    ASSERT_EQ(response.high_api_version().major(), 2);
    ASSERT_EQ(response.high_api_version().minor(), 12);
    ASSERT_EQ(response.high_api_version().patch(), 0);
    ASSERT_EQ(response.low_api_version().major(), 2);
    ASSERT_EQ(response.low_api_version().minor(), 1);
    ASSERT_EQ(response.low_api_version().patch(), 0);

    // Execution specific capabilities
    ASSERT_EQ(response.execution_capabilities().digest_functions_size(), 1);
#ifndef __SANITIZE_ADDRESS__
    ASSERT_EQ(response.execution_capabilities().digest_functions(0),
              DigestGenerator::digestFunction());
#endif
    ASSERT_EQ(response.execution_capabilities().digest_function(),
              DigestGenerator::digestFunction());
    ASSERT_TRUE(response.execution_capabilities().exec_enabled());
}

TEST_F(LocalExecutionServiceTestFixture, ExecuteExport)
{
    Action action;
    action.mutable_input_root_digest()->CopyFrom(DigestGenerator::hash(""));
    Command command;
    command.add_arguments("/bin/sh");
    command.add_arguments("-c");
    command.add_arguments("echo hello, world > hello.out");
    command.add_output_paths("hello.out");

    const auto commandDigest = casClient->uploadMessage(command);
    action.mutable_command_digest()->CopyFrom(commandDigest);
    const auto actionDigest = casClient->uploadMessage(action);

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

    std::vector<OutputFile> files(actionResult.output_files().begin(),
                                  actionResult.output_files().end());
    EXPECT_EQ(files.size(), 1);

    buildboxcommon::TemporaryDirectory tmpdir;
    auto exportResponse = casClient->exportFiles(files, tmpdir.strname());
    EXPECT_EQ(exportResponse.responses_size(), 1);
    for (const auto &resp : exportResponse.responses()) {
        EXPECT_EQ(resp.status().code(), GRPC_STATUS_OK);
    }

    EXPECT_EQ(
        FileUtils::getFileContents((tmpdir.strname() + "/hello.out").c_str()),
        "hello, world\n");
}
