/*
 * Copyright 2022 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 <atomic>
#include <cstdint>
#include <gmock/gmock.h>
#include <grpcpp/grpcpp.h>
#include <gtest/gtest.h>
#include <optional>
#include <utility>

#include <buildboxcommon_remoteexecutionclient.h>
#include <buildboxcommon_temporaryfile.h>
#include <trexe_executioncontext.h>
#include <trexe_executionoptions.h>
using namespace trexe;
using namespace buildboxcommon;

typedef std::shared_ptr<buildboxcommon::GrpcClient> GrpcClientPtr;

const testing::internal::AnythingMatcher Any = {};
using ::testing::DoAll;
using ::testing::NiceMock;
using ::testing::Return;
using ::testing::SetArgPointee;
using ::testing::Throw;

class MockREClient : public RemoteExecutionClient {
  public:
    MockREClient(GrpcClientPtr reBaseClient, GrpcClientPtr casBaseClient)
        : RemoteExecutionClient(std::move(reBaseClient),
                                std::move(casBaseClient))
    {
    }
    MOCK_METHOD0(init, void());
    MOCK_METHOD4(asyncExecuteAction,
                 google::longrunning::Operation(
                     const Digest &actionDigest,
                     const std::atomic_bool &stop_requested, bool skipCache,
                     const ExecutionPolicy *executionPolicy));
    MOCK_METHOD4(executeAction,
                 ActionResult(const Digest &actionDigest,
                              const std::atomic_bool &stop_requested,
                              bool skipCache,
                              const ExecutionPolicy *executionPolicy));
    MOCK_METHOD1(cancelOperation, bool(const std::string &operationName));
    MOCK_METHOD3(fetchFromActionCache,
                 bool(const Digest &actionDigest,
                      const std::set<std::string> &outputs,
                      ActionResult *result));
    MOCK_METHOD1(getOperation, google::longrunning::Operation(
                                   const std::string &operationName));
};
class MockCASClient : public CASClient {
  public:
    MockCASClient(GrpcClientPtr casBaseClient)
        : CASClient(std::move(casBaseClient))
    {
    }
    MOCK_METHOD1(init, void(bool checkCapabilities));
    MOCK_METHOD2(findMissingBlobs,
                 std::vector<Digest>(const std::vector<Digest> &digests,
                                     GrpcClient::RequestStats *requestStats));
    MOCK_METHOD2(uploadBlobs, std::vector<UploadResult>(
                                  const std::vector<UploadRequest> &requests,
                                  GrpcClient::RequestStats *requestStats));
    MOCK_METHOD2(fetchStringImpl,
                 std::string(const Digest &digest,
                             GrpcClient::RequestStats *requestStats));
    MOCK_METHOD(void, fetchStream,
                (const Digest &, const StreamConsumer &, int64_t,
                 std::optional<int64_t>,
                 GrpcClient::RequestStats *requestStats));
    std::string
    fetchString(const Digest &digest,
                GrpcClient::RequestStats *requestStats = nullptr) override
    {
        return fetchStringImpl(digest, requestStats);
    }
};

class ExecutionContextTest : public ::testing::Test {
  protected:
    std::shared_ptr<MockREClient> mockReClient;
    std::shared_ptr<MockCASClient> mockCasClient;
    std::shared_ptr<ExecutionOptions> options;

    ExecutionContextTest()
    {
        // Objects needed for initializing the mock clients.
        // Their code should not be exercised in the tests.
        GrpcClientPtr reBaseClient = std::make_shared<GrpcClient>();
        GrpcClientPtr casBaseClient = std::make_shared<GrpcClient>();

        // Mocks used for testing the ExecutionContext
        mockReClient = std::make_shared<NiceMock<MockREClient>>(reBaseClient,
                                                                casBaseClient);
        mockCasClient =
            std::make_shared<NiceMock<MockCASClient>>(casBaseClient);

        options = std::make_shared<ExecutionOptions>();
    }

    ExecutionContext makeContext()
    {
        return ExecutionContext(options, this->mockCasClient,
                                this->mockReClient);
    }

    void SetUp() override
    {
        ON_CALL(*(this->mockReClient), init()).WillByDefault(Return());
        ON_CALL(*(this->mockCasClient), init(true)).WillByDefault(Return());
    }
};

TEST_F(ExecutionContextTest, TestCancelOperation)
{
    // GIVEN
    options->d_cancelMode = true;
    options->d_operation = "myoperation";
    ExecutionContext context = makeContext();

    // THEN
    EXPECT_CALL(*(this->mockReClient), cancelOperation("myoperation"));

    // WHEN
    context.cancelOperation();
}

TEST_F(ExecutionContextTest, CancelAuthenticationGrpcError)
{
    // GIVEN
    options->d_cancelMode = true;
    options->d_operation = "myoperation";
    ExecutionContext context = makeContext();
    grpc::Status grpcError(grpc::StatusCode::UNAUTHENTICATED,
                           "Unauthenticated");
    EXPECT_CALL(*(this->mockReClient), cancelOperation(Any))
        .WillOnce(
            Throw(buildboxcommon::GrpcError("unauthenticated", grpcError)));
    // WHEN
    EXPECT_THROW(context.cancelOperation(), buildboxcommon::GrpcError);

    // THEN
    ASSERT_EQ(context.resultMetadata().error().code(),
              grpc::StatusCode::UNAUTHENTICATED);
}

class ExecutionContextDownloadOperationTest : public ExecutionContextTest {
  protected:
    Operation op;
    ActionResult actionResult;

    ExecutionContextDownloadOperationTest()
    {
        op.set_done(true);
        auto response = google::protobuf::Any();
        response.PackFrom(ExecuteResponse());
        *op.mutable_response() = response;
        actionResult.set_exit_code(0);
    }
};

TEST_F(ExecutionContextDownloadOperationTest, WithSuccess)
{
    // GIVEN
    EXPECT_CALL(*(this->mockReClient), getOperation(Any)).WillOnce(Return(op));
    auto context = makeContext();

    // WHEN
    bool success = context.downloadCompletedOperation("output_dir", false);

    // THEN
    ASSERT_TRUE(success);
}

TEST_F(ExecutionContextDownloadOperationTest, WithOperationError)
{
    // GIVEN
    auto status = google::rpc::Status();
    status.set_code(1);
    *op.mutable_error() = status;
    EXPECT_CALL(*(this->mockReClient), getOperation(Any)).WillOnce(Return(op));
    auto context = makeContext();

    // WHEN
    bool success = context.downloadCompletedOperation("output_dir", false);

    // THEN
    ASSERT_FALSE(success);
}

TEST_F(ExecutionContextDownloadOperationTest, NotDone)
{
    // GIVEN
    op.set_done(false);
    EXPECT_CALL(*(this->mockReClient), getOperation(Any)).WillOnce(Return(op));
    auto context = makeContext();

    // WHEN
    bool success = context.downloadCompletedOperation("output_dir", false);

    // THEN
    ASSERT_FALSE(success);
}

TEST_F(ExecutionContextDownloadOperationTest, InternalFailure)
{
    // GIVEN
    op.set_done(false);
    EXPECT_CALL(*(this->mockReClient), getOperation(Any))
        .WillOnce(Throw(std::runtime_error("service error")));
    auto context = makeContext();

    // WHEN
    EXPECT_THROW(context.downloadCompletedOperation("output_dir", false),
                 std::runtime_error);
}

TEST_F(ExecutionContextDownloadOperationTest, AuthenticationGrpcError)
{
    // GIVEN
    grpc::Status grpcError(grpc::StatusCode::UNAUTHENTICATED,
                           "Unauthenticated");
    EXPECT_CALL(*(this->mockReClient), getOperation(Any))
        .WillOnce(
            Throw(buildboxcommon::GrpcError("unauthenticated", grpcError)));
    auto context = makeContext();

    // WHEN
    EXPECT_THROW(context.downloadCompletedOperation("output_dir", false),
                 buildboxcommon::GrpcError);

    // THEN
    ASSERT_EQ(context.resultMetadata().error().code(),
              grpc::StatusCode::UNAUTHENTICATED);
}

class ExecutionContextResultStreamsTest : public ExecutionContextTest {
  protected:
    // Literals
    // Other members
    ActionData data;
    ActionResult returnedActionResult;
    TemporaryFile stdoutFile;
    TemporaryFile stderrFile;

    ExecutionContextResultStreamsTest()
    {
        // Set up ActionData
        Digest returnDigest;
        // (These values don't really matter)
        returnDigest.set_hash("myactionresultdigesthash");
        returnDigest.set_size_bytes(42);

        data.d_actionDigest = returnDigest;

        // Set up stream digests for ActionResult
        // (the values don't really matter)
        Digest stdoutDigest;
        Digest stderrDigest;
        stdoutDigest.set_hash("mystdoutdigesthash");
        stdoutDigest.set_size_bytes(10);
        stderrDigest.set_hash("my stderr digest hash");
        stderrDigest.set_size_bytes(20);

        // Set up ActionResult
        returnedActionResult.set_stdout_raw("Raw stdout");
        returnedActionResult.set_stderr_raw("Raw stderr");
        *returnedActionResult.mutable_stdout_digest() = stdoutDigest;
        *returnedActionResult.mutable_stderr_digest() = stderrDigest;

        // Configure options
        options->d_stdoutFile = stdoutFile.name();
        options->d_stderrFile = stderrFile.name();
    }
};

TEST_F(ExecutionContextResultStreamsTest, FalseIfNoActionResult)
{
    // GIVEN
    ExecutionContext context = makeContext();
    context.setActionData(data);

    // THEN
    EXPECT_FALSE(context.outputStdoutStderr());
}

TEST_F(ExecutionContextResultStreamsTest, ReadInlinedStream)
{
    // GIVEN
    ExecutionContext context = makeContext();
    context.setActionData(data);
    // The context has a populated ActionResult, i.e. it's done
    EXPECT_CALL(*(this->mockReClient), fetchFromActionCache(Any, Any, Any))
        .WillOnce(DoAll(SetArgPointee<2>(returnedActionResult), Return(true)));
    context.getActionResult(true);

    // THEN
    // We should not expect a call to fetch the stream contents from CAS
    EXPECT_CALL(*(this->mockCasClient), fetchStringImpl(Any, Any)).Times(0);

    // WHEN
    EXPECT_NO_THROW(context.outputStdoutStderr());

    // AND THEN
    // The returned contents should be the same as what's provided in the
    // ActionResult
    ASSERT_EQ(FileUtils::getFileContents(stdoutFile.fd()),
              returnedActionResult.stdout_raw());
    ASSERT_EQ(FileUtils::getFileContents(stderrFile.fd()),
              returnedActionResult.stderr_raw());
}

TEST_F(ExecutionContextResultStreamsTest, FetchFromCasIfEmptyInlined)
{
    // GIVEN
    // An empty inlined stdout
    returnedActionResult.clear_stdout_raw();
    returnedActionResult.clear_stderr_raw();
    std::string streamStdoutFromCAS = "Streamdata from CAS";
    std::string streamStderrFromCAS = "Stream stderr from CAS";

    ExecutionContext context = makeContext();
    context.setActionData(data);
    // The context has a populated ActionResult, i.e. it's done
    EXPECT_CALL(*(this->mockReClient), fetchFromActionCache(Any, Any, Any))
        .WillOnce(DoAll(SetArgPointee<2>(returnedActionResult), Return(true)));
    context.getActionResult(true);

    // THEN
    // The stream contents are fetched from CAS
    auto mockFetchStdoutStream = [&](const Digest &digest,
                                     const StreamConsumer &streamConsumer,
                                     int64_t, std::optional<int64_t>,
                                     GrpcClient::RequestStats *requestStats) {
        ReadResponse response;
        response.set_data(streamStdoutFromCAS);
        streamConsumer(response, 0);
    };
    auto mockFetchStderrStream = [&](const Digest &digest,
                                     const StreamConsumer &streamConsumer,
                                     int64_t, std::optional<int64_t>,
                                     GrpcClient::RequestStats *requestStats) {
        ReadResponse response;
        response.set_data(streamStderrFromCAS);
        streamConsumer(response, 0);
    };
    EXPECT_CALL(*(this->mockCasClient), fetchStream(Any, Any, Any, Any, Any))
        .WillOnce(testing::Invoke(mockFetchStdoutStream))
        .WillOnce(testing::Invoke(mockFetchStderrStream));

    // WHEN
    EXPECT_NO_THROW(context.outputStdoutStderr());

    // AND THEN
    ASSERT_EQ(FileUtils::getFileContents(stdoutFile.fd()),
              streamStdoutFromCAS);
    ASSERT_EQ(FileUtils::getFileContents(stderrFile.fd()),
              streamStderrFromCAS);
}
