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

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

/***
 * Fixture to support testing buildbox-casd as a hybrid Execution service
 */
class HybridExecutionServiceTestFixture : public ::testing::Test {
  protected:
    buildboxcommon::TemporaryDirectory socketDirectory;
    const std::string TEST_REMOTE_SERVER_ADDRESS;
    const std::string TEST_PROXY_SERVER_ADDRESS;

    buildboxcommon::TemporaryDirectory remoteStorageRootDirectory;
    std::shared_ptr<FsLocalCas> remoteCasStorage;
    std::shared_ptr<FsLocalAssetStorage> remoteAssetStorage;
    std::shared_ptr<FsLocalActionStorage> remoteActionStorage;
    std::shared_ptr<Server> remoteServer;

    buildboxcommon::TemporaryDirectory proxyStorageRootDirectory;
    std::shared_ptr<FsLocalCas> proxyCasStorage;
    std::shared_ptr<FsLocalAssetStorage> proxyAssetStorage;
    std::shared_ptr<FsLocalActionStorage> proxyActionStorage;
    std::shared_ptr<Server> proxyServer;

    std::shared_ptr<grpc::Channel> channel;
    std::shared_ptr<ByteStream::StubInterface> byteStreamStub;
    std::shared_ptr<ContentAddressableStorage::StubInterface> casStub;
    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 remoteInstanceName = "testremote";
    const std::string proxyInstanceName = "testproxy";

    HybridExecutionServiceTestFixture()
        : TEST_REMOTE_SERVER_ADDRESS("unix://" +
                                     std::string(socketDirectory.name()) +
                                     "/remote.sock"),
          TEST_PROXY_SERVER_ADDRESS(
              "unix://" + std::string(socketDirectory.name()) + "/proxy.sock"),
          grpcClient(std::make_shared<buildboxcommon::GrpcClient>()),
          casClient(std::make_shared<CASClient>(grpcClient)),
          execClient(
              std::make_shared<RemoteExecutionClient>(grpcClient, nullptr))
    {
        const char *runnerCommand = getenv("BUILDBOX_RUN");
        EXPECT_NE(runnerCommand, nullptr);
        const std::vector<std::string> extraRunArgs = {};

        // Create a server instance to act as remote server
        remoteCasStorage =
            std::make_shared<FsLocalCas>(remoteStorageRootDirectory.name());
        remoteAssetStorage = std::make_shared<FsLocalAssetStorage>(
            remoteStorageRootDirectory.name());
        remoteActionStorage = std::make_shared<FsLocalActionStorage>(
            remoteStorageRootDirectory.name());

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

        remoteServer =
            std::make_shared<Server>(remoteCasStorage, remoteAssetStorage,
                                     remoteActionStorage, remoteScheduler);
        remoteServer->addLocalServerInstance(remoteInstanceName);
        remoteServer->addListeningPort(TEST_REMOTE_SERVER_ADDRESS);
        remoteServer->start();

        // Create a server instance to act as proxy to the remote server,
        // also supporting local execution.
        ConnectionOptions remoteServerEndpoint;
        remoteServerEndpoint.setInstanceName(remoteInstanceName);
        remoteServerEndpoint.setUrl(TEST_REMOTE_SERVER_ADDRESS);

        proxyCasStorage =
            std::make_shared<FsLocalCas>(proxyStorageRootDirectory.name());
        proxyAssetStorage = std::make_shared<FsLocalAssetStorage>(
            proxyStorageRootDirectory.name());
        proxyActionStorage = std::make_shared<FsLocalActionStorage>(
            proxyStorageRootDirectory.name());

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

        proxyServer =
            std::make_shared<Server>(proxyCasStorage, proxyAssetStorage,
                                     proxyActionStorage, localScheduler);
        proxyServer->addProxyInstance(
            proxyInstanceName, remoteServerEndpoint, remoteServerEndpoint,
            remoteServerEndpoint, remoteServerEndpoint,
            /* readOnlyRemote */ false, /* proxyFindmissingblobsCacheTtl */ 0,
            /* proxyAcValidationCacheTtl */ 0,
            /* execHybridQueueLimit */ 0);
        proxyServer->addListeningPort(TEST_PROXY_SERVER_ADDRESS);
        proxyServer->start();

        // Initialize a gRPC client to the proxy
        channel = grpc::CreateChannel(TEST_PROXY_SERVER_ADDRESS,
                                      grpc::InsecureChannelCredentials());
        byteStreamStub = ByteStream::NewStub(channel);
        casStub = ContentAddressableStorage::NewStub(channel);
        executionStub = Execution::NewStub(channel);
        operationsStub = Operations::NewStub(channel);
        grpcClient->setInstanceName(proxyInstanceName);
        casClient->init(byteStreamStub, casStub, nullptr, 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);

        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);
    }

    void expectOperationSuccessful(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_EQ(actionResult.exit_code(), 0);
    }

    bool isRemoteOperation(const Operation &operation)
    {
        const std::string remotePrefix =
            proxyInstanceName + "#" + remoteInstanceName + "#";
        return operation.name().substr(0, remotePrefix.size()) == remotePrefix;
    }
};

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

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

TEST_F(HybridExecutionServiceTestFixture, 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(HybridExecutionServiceTestFixture, 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(HybridExecutionServiceTestFixture, CancelOperation)
{
    const auto start = std::chrono::steady_clock::now();

    const auto actionDigest = generateSimpleShellAction("sleep 10");

    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);

    const auto end = std::chrono::steady_clock::now();
    EXPECT_LT(end - start, std::chrono::seconds(3));
}

TEST_F(HybridExecutionServiceTestFixture, Queue)
{
    const auto start = std::chrono::steady_clock::now();

    // Test servers are configured for 2 parallel jobs each (local + remote)

    const auto actionDigest1 = generateSimpleShellAction("sleep 2 && echo 1");
    const auto actionDigest2 = generateSimpleShellAction("sleep 2 && echo 2");
    const auto actionDigest3 = generateSimpleShellAction("sleep 2 && echo 3");
    const auto actionDigest4 = generateSimpleShellAction("echo hello, world");
    const auto actionDigest5 = generateSimpleShellAction("sleep 10");

    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);
    auto operation5 =
        execClient->asyncExecuteAction(actionDigest5, stopRequested);

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

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

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

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

    // The three `sleep 2` should be executed in parallel (two locally and one
    // remotely). The total duration will be slightly longer than 2s but
    // verify it's below 4s to confirm the parallel execution and that
    // `sleep 10` was cancelled.
    EXPECT_GT(end - start, std::chrono::seconds(2));
    EXPECT_LT(end - start, std::chrono::seconds(4));

    EXPECT_FALSE(isRemoteOperation(operation1));
    EXPECT_FALSE(isRemoteOperation(operation2));
    EXPECT_TRUE(isRemoteOperation(operation3));
    EXPECT_TRUE(isRemoteOperation(operation4));
    EXPECT_TRUE(isRemoteOperation(operation5));
}

TEST_F(HybridExecutionServiceTestFixture, 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());
}
